Merge branch 'pr/4'

.gitignore

@@ -1,4 +1,5 @@
 venv
 .env
 __pycache__
-.vscode
+.vscode
+corpus

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI
 from langgraph.agents.summarize_agent.graph import graph   
+from langgraph.agents.rag_agent.graph import graph as rag_graph
 from fastapi import Request
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -31,6 +32,23 @@ async def summarize(request: Request):
     notes = data.get("notes")
     return  graph.invoke({"urls": urls, "codes": codes, "notes": notes})
 
+@app.post("/chat")
+async def chat(request: Request):
+    data = await request.json()
+    user_input = data.get("message", "")
+    chat_history = data.get("chat_history", [])
+    
+    # Invoke the RAG chatbot graph
+    result = rag_graph.invoke({
+        "user_input": user_input,
+        "chat_history": chat_history
+    })
+    
+    return {
+        "response": result["response"],
+        "chat_history": result["chat_history"]
+    }
+
 
 
 

langgraph/agents/rag_agent/graph.py

@@ -0,0 +1,207 @@
+import os
+from typing import Dict, List, Any, Literal
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langgraph.graph import StateGraph
+from langgraph.graph.graph import END
+from dotenv import load_dotenv
+import google.generativeai as genai
+from google.generativeai import GenerativeModel
+import sys
+
+# Add the parent directory to the path to import utils
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+from utils.create_vectordb import query_chroma_db
+
+load_dotenv()
+
+# Initialize Gemini model
+api_key = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=api_key)
+model = GenerativeModel("gemini-2.5-flash-preview-05-20")
+
+def retrieve_context(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Retrieve relevant context from the vector database based on the user query.
+    """
+    query = state.get("user_input", "")
+    if not query:
+        return {"context": "No query provided.", "user_input": query, "next": "request_clarification"}
+    
+    # Check if query is clear enough
+    if len(query.split()) < 3 or "?" not in query and not any(w in query.lower() for w in ["what", "how", "why", "when", "where", "who", "which"]):
+        return {"context": "", "user_input": query, "next": "request_clarification"}
+    
+    # Query the vector database
+    results = query_chroma_db(query, n_results=3)
+    
+    # Extract the retrieved documents
+    documents = results.get("documents", [[]])[0]
+    metadatas = results.get("metadatas", [[]])[0]
+    
+    # Format the context
+    formatted_context = []
+    for i, (doc, metadata) in enumerate(zip(documents, metadatas)):
+        source = metadata.get("source", "Unknown")
+        formatted_context.append(f"Document {i+1} (Source: {source}):\n{doc}\n")
+    
+    context = "\n".join(formatted_context) if formatted_context else ""
+    
+    # Determine next step based on context quality
+    if not context or len(context) < 50:
+        next_step = "use_gemini_knowledge"
+    else:
+        next_step = "generate_response"
+    
+    return {"context": context, "user_input": query, "next": next_step}
+
+def request_clarification(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Request clarification from the user when the query is unclear.
+    """
+    query = state.get("user_input", "")
+    
+    clarification_message = model.generate_content(
+        f"""The user asked: "{query}"
+        
+        This query seems vague or unclear. Generate a polite response asking for more specific details.
+        Focus on what additional information would help you understand their request better.
+        Keep your response under 3 sentences and make it conversational."""
+    )
+    
+    response = clarification_message.text
+    
+    # Update chat history
+    chat_history = state.get("chat_history", [])
+    new_chat_history = chat_history + [
+        {"role": "user", "content": query},
+        {"role": "assistant", "content": response}
+    ]
+    
+    
+    return {
+        "response": response,
+        "chat_history": new_chat_history,
+        "needs_clarification": True
+    }
+
+def use_gemini_knowledge(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Use Gemini's knowledge base when no relevant information is found in the vector database.
+    """
+    query = state.get("user_input", "")
+    chat_history = state.get("chat_history", [])
+    
+    # Construct the prompt
+    prompt_template = """
+    I couldn't find specific information about this in my local database. However, I can try to answer based on my general knowledge.
+    
+    User Question: {query}
+    
+    First, acknowledge that you're answering from general knowledge rather than the specific database.
+    Then provide a helpful, accurate response based on what you know about the topic.
+    """
+    
+    # Generate response
+    response = model.generate_content(
+        prompt_template.format(query=query)
+    )
+    
+    # Update chat history
+    new_chat_history = chat_history + [
+        {"role": "user", "content": query},
+        {"role": "assistant", "content": response.text}
+    ]
+    
+    return {
+        "response": response.text,
+        "chat_history": new_chat_history
+    }
+
+def generate_response(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Generate a response using the LLM based on the retrieved context and user query.
+    """
+    context = state.get("context", "")
+    query = state.get("user_input", "")
+    chat_history = state.get("chat_history", [])
+    
+    # Construct the prompt
+    prompt_template = """
+    You are a helpful assistant that answers questions based on the provided context.
+    
+    Context:
+    {context}
+    
+    Chat History:
+    {chat_history}
+    
+    User Question: {query}
+    
+    Answer the question based only on the provided context. If the context doesn't contain enough information,
+    acknowledge this but still try to provide a helpful response based on the available information.
+    Provide a clear, concise, and helpful response.
+    """
+    
+    # Format chat history for the prompt
+    formatted_chat_history = "\n".join([f"{msg['role']}: {msg['content']}" for msg in chat_history])
+    
+    # Generate response
+    response = model.generate_content(
+        prompt_template.format(
+            context=context,
+            chat_history=formatted_chat_history,
+            query=query
+        )
+    )
+    
+    # Update chat history
+    new_chat_history = chat_history + [
+        {"role": "user", "content": query},
+        {"role": "assistant", "content": response.text}
+    ]
+    
+    return {
+        "response": response.text,
+        "chat_history": new_chat_history
+    }
+
+def decide_next_step(state: Dict[str, Any]) -> Literal["request_clarification", "use_gemini_knowledge", "generate_response"]:
+    """
+    Decide the next step in the workflow based on the state.
+    """
+    return state["next"]
+
+# Define the workflow
+def build_graph():
+    workflow = StateGraph(state_schema=Dict[str, Any])
+    
+    # Add nodes
+    workflow.add_node("retrieve_context", retrieve_context)
+    workflow.add_node("request_clarification", request_clarification)
+    workflow.add_node("use_gemini_knowledge", use_gemini_knowledge)
+    workflow.add_node("generate_response", generate_response)
+    
+    # Define edges with conditional routing
+    workflow.set_entry_point("retrieve_context")
+    workflow.add_conditional_edges(
+        "retrieve_context",
+        decide_next_step,
+        {
+            "request_clarification": "request_clarification",
+            "use_gemini_knowledge": "use_gemini_knowledge",
+            "generate_response": "generate_response"
+        }
+    )
+    
+    # Set finish points
+    workflow.add_edge("request_clarification", END)
+    workflow.add_edge("use_gemini_knowledge", END)
+    workflow.add_edge("generate_response", END)
+    
+    # Compile the graph
+    return workflow.compile()
+
+# Create the graph
+graph = build_graph()

requirements.txt

@@ -3,6 +3,10 @@ uvicorn[standard]
 langgraph 
 langsmith
 google-genai
-
+google-generativeai
+chromadb
+langchain
+langchain-community
 python-dotenv
+pypdf
 

utils/__init__.py

@@ -0,0 +1 @@
+# This file is intentionally left empty to make the directory a Python package

utils/create_vectordb.py

@@ -0,0 +1,141 @@
+import os
+from typing import Optional, List
+import chromadb
+from chromadb.utils import embedding_functions
+from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from dotenv import load_dotenv
+import google.generativeai as genai
+
+load_dotenv()
+
+# Configure paths
+CORPUS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "corpus")
+DB_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "vectordb")
+
+# Ensure directories exist
+os.makedirs(CORPUS_DIR, exist_ok=True)
+os.makedirs(DB_DIR, exist_ok=True)
+
+def load_documents(corpus_dir: str = CORPUS_DIR) -> List:
+    """Load documents from the corpus directory."""
+    if not os.path.exists(corpus_dir):
+        raise FileNotFoundError(f"Corpus directory not found: {corpus_dir}")
+    print(f"Loading documents from {corpus_dir}...")
+    
+    # Initialize loaders for different file types
+    loaders = {
+        # "txt": DirectoryLoader(corpus_dir, glob="**/*.txt", loader_cls=TextLoader),
+        "pdf": DirectoryLoader(corpus_dir, glob="**/*.pdf", loader_cls=PyPDFLoader),
+        # "docx": DirectoryLoader(corpus_dir, glob="**/*.docx", loader_cls=Docx2txtLoader),
+    }
+    
+    documents = []
+    for file_type, loader in loaders.items():
+        try:
+            docs = loader.load()
+            print(f"Loaded {len(docs)} {file_type} documents")
+            documents.extend(docs)
+        except Exception as e:
+            print(f"Error loading {file_type} documents: {e}")
+    
+    return documents
+
+def split_documents(documents, chunk_size=1000, chunk_overlap=200):
+    """Split documents into chunks."""
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len,
+    )
+    
+    splits = text_splitter.split_documents(documents)
+    print(f"Split {len(documents)} documents into {len(splits)} chunks")
+    
+    return splits
+
+def create_chroma_db(documents, collection_name="corpus_collection", db_dir=DB_DIR):
+    """Create a Chroma vector database from documents."""
+    # Initialize the Gemini embedding function
+    gemini_ef = embedding_functions.GoogleGenerativeAiEmbeddingFunction(
+        api_key=os.getenv("GOOGLE_API_KEY"),
+        model_name="models/embedding-001"
+    )
+    
+    # Initialize Chroma client
+    client = chromadb.PersistentClient(path=db_dir)
+    
+    # Create or get collection
+    try:
+        collection = client.get_collection(name=collection_name)
+        print(f"Using existing collection: {collection_name}")
+    except:
+        collection = client.create_collection(
+            name=collection_name,
+            embedding_function=gemini_ef
+        )
+        print(f"Created new collection: {collection_name}")
+    
+    # Add documents to collection
+    for i, doc in enumerate(documents):
+        collection.add(
+            documents=[doc.page_content],
+            metadatas=[doc.metadata],
+            ids=[f"doc_{i}"]
+        )
+    
+    print(f"Added {len(documents)} documents to collection {collection_name}")
+    return collection
+
+def query_chroma_db(query: str, collection_name="corpus_collection", n_results=5, db_dir=DB_DIR):
+    """Query the Chroma vector database."""
+    # Initialize the Gemini embedding function
+    gemini_ef = embedding_functions.GoogleGenerativeAiEmbeddingFunction(
+        api_key=os.getenv("GOOGLE_API_KEY"),
+        model_name="models/embedding-001"
+    )
+    
+    # Initialize Chroma client
+    client = chromadb.PersistentClient(path=db_dir)
+    
+    # Get collection
+    collection = client.get_collection(name=collection_name, embedding_function=gemini_ef)
+    
+    # Query collection
+    results = collection.query(
+        query_texts=[query],
+        n_results=n_results
+    )
+    
+    return results
+
+def main():
+    """Main function to create and test the vector database."""
+    print("Starting vector database creation...")
+    
+    # Load documents
+    documents = load_documents()
+    if not documents:
+        print("No documents found in corpus directory. Please add documents to proceed.")
+        return
+    
+    # Split documents
+    splits = split_documents(documents)
+    
+    # Create vector database
+    collection = create_chroma_db(splits)
+    
+    # Test query
+    test_query = "What is this corpus about?"
+    print(f"\nTesting query: '{test_query}'")
+    results = query_chroma_db(test_query)
+    print(f"Found {len(results['documents'][0])} matching documents")
+    for i, (doc, metadata) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
+        print(f"\nResult {i+1}:")
+        print(f"Document: {doc[:150]}...")
+        print(f"Source: {metadata.get('source', 'Unknown')}")
+    
+    print("\nVector database creation and testing complete!")
+
+if __name__ == "__main__":
+    main()