shovo896/codedocumentation

1import os2 3from dotenv import load_dotenv4from langchain_core.prompts import ChatPromptTemplate5from langchain_openai import ChatOpenAI, OpenAIEmbeddings6from langchain_pinecone import PineconeVectorStore7 8load_dotenv()9 10INDEX_NAME = "code-doc-search-openai"11EMBED_MODEL = "text-embedding-3-small"12EMBED_DIM = 153613 14 15def get_required_env(name: str) -> str:16    value = os.getenv(name, "").strip().strip('"').strip("'")17    if not value:18        raise RuntimeError(f"Missing required environment variable: {name}")19    return value20 21 22def get_qa_chain(namespace=None, repo_url=None, branch=None):23    embeddings = OpenAIEmbeddings(24        model=EMBED_MODEL,25        dimensions=EMBED_DIM,26        api_key=get_required_env("OPENAI_API_KEY"),27    )28 29    vectorstore = PineconeVectorStore(30        index_name=INDEX_NAME,31        embedding=embeddings,32        namespace=namespace,33    )34 35    llm = ChatOpenAI(36        model_name="gpt-4o-mini",37        api_key=get_required_env("OPENAI_API_KEY"),38        temperature=0.2,39    )40 41    repo_context = ""42    if repo_url:43        repo_context = f"\nRepository: {repo_url}"44        if branch:45            repo_context += f"\nBranch: {branch}"46 47    system_prompt = f"""You are a helpful assistant for answering questions about code documentation.48Always answer in English, even if the user's question is written in another language.49Use only the following retrieved repository context to answer the user's question.50If the answer is not in the retrieved context, say "I don't know."51When asked about technologies or frameworks, infer from filenames, file extensions, package manifests, config files, imports, and scripts in the retrieved context.52Treat .ts and .tsx files as TypeScript evidence, and distinguish that from plain .js or .jsx JavaScript files.53For yes/no technology questions, give the strongest supported conclusion from the evidence instead of saying "I don't know" when file extensions, dependencies, or config files are enough to answer.54A repo with package.json, Vite, React, tsconfig.json, .ts, or .tsx files is a JavaScript/TypeScript frontend project; if only TypeScript evidence is present, say it uses TypeScript rather than plain JavaScript.55{repo_context}56 57Context:58{{context}}"""59 60    prompt = ChatPromptTemplate.from_messages(61        [62            ("system", system_prompt),63            ("human", "{query}"),64        ]65    )66 67    normalization_prompt = ChatPromptTemplate.from_messages(68        [69            (70                "system",71                "Rewrite the user's question as a concise English question for codebase search. "72                "The user may write in Bengali, romanized Bengali/Banglish, Hindi, or mixed English. "73                "Preserve technical terms, package names, filenames, and framework names. "74                "Return only the rewritten English question.",75            ),76            ("human", "{query}"),77        ]78    )79 80    retriever = vectorstore.as_retriever(search_kwargs={"k": 8})81 82    def format_docs(docs):83        formatted = []84        for doc in docs:85            source = doc.metadata.get("source", "unknown")86            file_type = doc.metadata.get("file_type", "")87            formatted.append(88                f"Source: {source}\n"89                f"File type: {file_type}\n"90                f"Content:\n{doc.page_content}"91            )92        return "\n\n---\n\n".join(formatted)93 94    chain = prompt | llm95    normalizer = normalization_prompt | llm96 97    class QAChain:98        def __init__(self, retriever, llm_chain, normalize_chain):99            self.retriever = retriever100            self.llm_chain = llm_chain101            self.normalize_chain = normalize_chain102 103        def invoke(self, inputs):104            original_query = inputs.get("query") if isinstance(inputs, dict) else inputs105            query = original_query106            try:107                normalized = self.normalize_chain.invoke({"query": original_query})108                query = normalized.content.strip() or original_query109            except Exception:110                query = original_query111 112            docs = self.retriever.invoke(query)113            result = self.llm_chain.invoke(114                {115                    "context": format_docs(docs),116                    "query": query,117                }118            )119            return {120                "result": result.content,121                "source_documents": docs,122            }123 124    return QAChain(retriever, chain, normalizer)125