File size: 1,196 Bytes
d3dba58
26fe9e2
 
 
d3dba58
3bc9c63
dd452e0
26fe9e2
62f1a72
2c837a2
db5635e
2c837a2
 
3bc9c63
26fe9e2
 
dd452e0
 
 
 
 
 
 
26fe9e2
 
 
 
 
 
dd452e0
 
26fe9e2
3bc9c63
26fe9e2
 
dd452e0
 
 
 
26fe9e2
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

def build_rag_chain(vectorstore, groq_api_key):
    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

    llm = ChatGroq(
        api_key=groq_api_key,
        model="llama-3.1-8b-instant",
        temperature=0
    )

    prompt = ChatPromptTemplate.from_template(
        """
        You are a codebase assistant.

        RULES:
        - Use ONLY the context below.
        - If information is missing, say:
          "I cannot find this information in the provided codebase."
        - Do NOT guess.

        Context:
        {context}

        Question:
        {question}

        Answer:
        """
    )

    def format_docs(docs):
        return "\n\n".join(
            f"FILE: {d.metadata['file']}\n{d.page_content}"
            for d in docs
        )

    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough()
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain