Spaces:
Sleeping
Sleeping
File size: 2,800 Bytes
9806c71 |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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() |