Spaces:
Sleeping
Sleeping
| from typing import Annotated, List, TypedDict | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from langgraph.graph import StateGraph, START, END | |
| from langgraph.prebuilt import ToolNode | |
| from src.tools import tools | |
| # The Local Knowledge Registry (Update this whenever you add new data types) | |
| # Manual as of now | |
| LOCAL_MANIFEST = { | |
| "topics": ["HR Policies", "Project X Design Docs", "Q3 Financials", "Employee Handbook"], | |
| "date_range": "Documents updated as of Dec 2024", | |
| "domain": "Internal Corporate Knowledge" | |
| } | |
| SYSTEM_PROMPT = f""" | |
| You are an expert Research Assistant. You have access to: | |
| 1. INTERNAL DATA: {LOCAL_MANIFEST['topics']}. (Use 'local_research_tool') | |
| 2. EXTERNAL DATA: The entire internet via duckduckgosearch. (Use 'web_search_tool') | |
| GUIDELINES: | |
| - Given the user's technical question and the fact that our internal documents are insufficient, generate a generic search query for the internet that does NOT include any proprietary names or internal details. | |
| - If a query is about {LOCAL_MANIFEST['topics']}, try LOCAL first. | |
| - If a query is TECHNICAL (e.g., PyTorch, Python APIs) or REAL-TIME, go to WEB immediately. | |
| - If the query is ambiguous, try LOCAL first, then fallback to WEB if the results are empty or low confidence. | |
| """ | |
| class AgentState(TypedDict): | |
| """MNC Agent state with intent tracking and sufficiency grading.""" | |
| messages: Annotated[List, "Chat history"] | |
| intent: str | |
| is_sufficient: bool | |
| # Brain: Gemini 2.0 Flash (high-speed agentic reasoning) | |
| llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0) | |
| llm_with_tools = llm.bind_tools(tools) | |
| def router(state: AgentState): | |
| """Classifies user intent to prioritize retrieval paths.""" | |
| query = state['messages'][-1].content | |
| prompt = f"Categorize intent: TECHNICAL (API/Docs), INTERNAL (Proprietary), or REALTIME. Query: {query}" | |
| response = llm.invoke(prompt) | |
| intent = "TECHNICAL" if any(x in response.content.upper() for x in ["TECHNICAL", "REALTIME"]) else "INTERNAL" | |
| return {"intent": intent} | |
| def call_model(state: AgentState): | |
| """Invokes Gemini with tools based on intent and history.""" | |
| return {"messages": [llm_with_tools.invoke(state['messages'])]} | |
| # Orchestration Graph | |
| workflow = StateGraph(AgentState) | |
| workflow.add_node("router", router) | |
| workflow.add_node("llm", call_model) | |
| workflow.add_node("tools", ToolNode(tools)) | |
| workflow.add_edge(START, "router") | |
| workflow.add_edge("router", "llm") | |
| # Self-Correction Loop | |
| def should_continue(state: AgentState): | |
| last_msg = state["messages"][-1] | |
| return "tools" if last_msg.tool_calls else END | |
| workflow.add_conditional_edges("llm", should_continue) | |
| workflow.add_edge("tools", "llm") | |
| app = workflow.compile() |