rag.py

import re

from fastapi import APIRouter
from pydantic import BaseModel
from typing import List

from api.datastores.chroma import VectorStore
from api.plugins.llms import manager


router = APIRouter()


class RagPromptRequest(BaseModel):
    query: str
    memory: List[str]
    provider: str


@router.post("/generate")
def generate_rag_prompt(request: RagPromptRequest):
    llm = manager.LLMManager().get_provider(request.provider)()
    vector_store = VectorStore(collection_name="forensics")
    vector_query = request.query + " ".join(request.memory)
    response = vector_store.query_collection(query=vector_query)

    documents = [document for document in response.get("documents", [])[0]]
    metadata = [source["source"] for source in response.get("metadatas", [])[0]]
    distances = [distance for distance in response.get("distances", [])[0]]

    context = []
    for document in documents:
        context.append(f"<article>{document}</article>")
    prompt_template = """
    QUESTION:
    {query}

    INSTRUCTIONS:
    * You are a security engineer with 10 years of experience.
    * You will be provided a list of documents delimited with XML tags.
    * Answer the QUESTION using ONLY the DOCUMENTS text above.
    * Keep your answer ground in the facts of the DOCUMENT.
    * If the DOCUMENTS doesn't contain the facts to answer the QUESTION return "I can't answer that".
    * Output format MUST be in markdown.

    DOCUMENTS:
    {context}
    """
    prompt = llm.prompt_from_template(
        prompt_template, dict(context=context, query=request.query)
    )

    # Extract the source links.
    sources = list()
    unique_sources = set()
    for index, url in enumerate(metadata):
        if url in unique_sources:
            continue
        sources.append({"source": url, "distance": distances[index]})
        unique_sources.add(url)

    summary = llm.generate(prompt)

    # Generate suggested follow-up questions
    suggestions = get_suggested_questions(llm, summary, context)

    return {
        "prompt": prompt,
        "summary": summary,
        "sources": sources,
        "suggestions": suggestions,
    }


def get_suggested_questions(llm: object, summary: str, context: str) -> list:
    prompt_template = """
    INSTRUCTIONS:
    * I'm a security engineer investigating and I'm doing research on forensics.
    * You are a security minded research expert with 20 years of experience.
    * You will be provided a SUMMARY and DOCUMENTS.
    * Create MAX five suggested SHORT questions from the SUMMARY.
    * You MUST be able to answer the suggested questions using ONLY the provided DOCUMENTS.
    * Include important keywords from the SUMMARY in your suggestions.
    * Your response should be a numbered list with each item on a new line. For example: \n\n1. foo\n\n2. bar\n\n3. baz"

    SUMMARY:
    {summary}

    DOCUMENTS:
    {context}
    """
    prompt = llm.prompt_from_template(
        prompt_template, dict(summary=summary, context=context)
    )
    generated_suggestions = llm.generate(prompt)
    pattern = r"\d+\.\s([^\n]+)"
    suggestions = re.findall(pattern, generated_suggestions)
    return suggestions