Upload 8 files

agents/rag_agent.py

@@ -0,0 +1,75 @@
+from typing import Dict, Any, List
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema import Document
+from pydantic import BaseModel, Field
+from models.vector_store import VectorStore
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+class RAGAgentState(BaseModel):
+    """State for the RAG agent"""
+    query: str = Field(description="The user's query for document retrieval")
+    context: List[Dict[str, Any]] = Field(description="Retrieved context from documents", default=[])
+    response: str = Field(description="The agent's response", default="")
+
+class RAGAgent:
+    """Agent that handles document-based queries using RAG"""
+    
+    def __init__(self, api_key: str = GEMINI_API_KEY):
+        self.llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            google_api_key=api_key
+        )
+        
+        self.vector_store = VectorStore()
+        
+        self.rag_prompt = ChatPromptTemplate.from_messages([
+            ("system", """You are an expert research assistant helping a user understand complex topics clearly and concisely.
+                Use only the provided context to answer the user's question. If the context does not contain the answer, say:
+                "I don't have enough information to answer that question."
+
+                When answering:
+                - Explain technical terms simply
+                - Use examples if helpful
+                - Keep the tone friendly and helpful
+
+                Context:
+                {context}"""),
+                ("human", "{query}")
+            ])
+        
+        self.rag_chain = self.rag_prompt | self.llm
+    
+    def retrieve_context(self, query: str, k: int = 4) -> List[Document]:
+        """Retrieve relevant context from the vector store"""
+        return self.vector_store.similarity_search(query, k=k)
+    
+    def get_rag_response(self, query: str) -> Dict[str, Any]:
+        """Generate a RAG-based response to the query"""
+        # Retrieve relevant documents
+        docs = self.retrieve_context(query)
+        
+        if not docs:
+            return {
+                "context": [],
+                "response": "I couldn't find any relevant information in the documents to answer your question."
+            }
+        
+        # Format context
+        context_texts = [doc.page_content for doc in docs]
+        context_str = "\n\n".join(context_texts)
+        
+        # Generate response
+        response = self.rag_chain.invoke({
+            "query": query,
+            "context": context_str
+        })
+        
+        return {
+            "context": [{"page_content": doc.page_content, "metadata": doc.metadata} for doc in docs],
+            "response": response.content
+        }

agents/router_agent.py

@@ -0,0 +1,50 @@
+from typing import Dict, Any, List, Tuple
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.prompts import ChatPromptTemplate
+from pydantic import BaseModel , Field
+from langgraph.graph import StateGraph
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+class RouterState(BaseModel):
+    """State for the router agent"""
+    query: str = Field(description="The user's query")
+    action: str = Field(description="The action to take: 'weather' or 'document'")
+    
+class RouterAgent:
+    """Agent that decides whether to use weather API or document RAG"""
+    
+    def __init__(self, api_key: str = GEMINI_API_KEY):
+        self.llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            google_api_key=api_key
+        )
+        
+        self.prompt = ChatPromptTemplate.from_messages([
+            ("system", """You are a router agent that decides whether a user query is about:
+            1. Weather information (requiring a weather API call)
+            2. Information from documents (requiring RAG retrieval)
+            
+            If the query mentions weather, forecast, temperature, rain, sun, climate, or other weather-related terms for a specific location, classify it as 'weather'.
+            
+            Otherwise, classify it as 'document' for document retrieval.
+            
+            Return only 'weather' or 'document' as your classification."""),
+            ("human", "{query}")
+        ])
+        
+        self.chain = self.prompt | self.llm
+    
+    def route_query(self, query: str) -> str:
+        """Route a query to either weather API or document RAG"""
+        response = self.chain.invoke({"query": query})
+        # Extract just the decision: 'weather' or 'document'
+        decision = response.content.strip().lower()
+        
+        if "weather" in decision:
+            return "weather"
+        else:
+            return "document"

agents/weather_agent.py

@@ -0,0 +1,82 @@
+from typing import Dict, Any
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel, Field
+from utils.api_handler import WeatherAPIHandler
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+
+
+class WeatherAgentState(BaseModel):
+    """State for the weather agent"""
+    query: str = Field(description="The user's query about weather")
+    city: str = Field(description="The city to get weather for", default="")
+    weather_data: Dict[str, Any] = Field(description="Raw weather data", default={})
+    response: str = Field(description="The agent's response", default="")
+
+class WeatherAgent:
+    """Agent that handles weather-related queries"""
+    
+    def __init__(self, api_key: str = GEMINI_API_KEY):
+        self.llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            google_api_key=api_key
+        )
+        
+        self.weather_api = WeatherAPIHandler()
+        
+        self.extract_city_prompt = ChatPromptTemplate.from_messages([
+            ("system", """Extract the city name from the user's weather query.
+            Return ONLY the city name, nothing else.
+            If no city is mentioned, return "Not specified"."""),
+            ("human", "{query}")
+        ])
+        
+        self.response_prompt = ChatPromptTemplate.from_messages([
+            ("system", """You are a helpful weather assistant.
+            Format the weather information in a friendly, conversational way.
+            Include all relevant weather details from the provided data."""),
+            ("human", "Query: {query}\nWeather Data: {weather_info}")
+        ])
+        
+        self.extract_city_chain = self.extract_city_prompt | self.llm
+        self.response_chain = self.response_prompt | self.llm
+    
+    def extract_city(self, query: str) -> str:
+        """Extract city name from the user query"""
+        response = self.extract_city_chain.invoke({"query": query})
+        city = response.content.strip()
+        
+        # Handle case where no city is specified
+        if city.lower() == "not specified":
+            return "London"  # Default city
+        
+        return city
+    
+    def get_weather_response(self, query: str, city: str = None) -> Dict[str, Any]:
+        """Get weather data and generate a response"""
+        # Extract city if not provided
+        if not city:
+            city = self.extract_city(query)
+        
+        # Get weather data
+        weather_data = self.weather_api.get_weather(city)
+        
+        # Format weather data
+        weather_info = self.weather_api.format_weather_data(weather_data)
+        
+        # Generate response
+        response = self.response_chain.invoke({
+            "query": query,
+            "weather_info": weather_info
+        })
+        
+        return {
+            "city": city,
+            "weather_data": weather_data,
+            "response": response.content
+        }