Create agent.py

agent.py

@@ -0,0 +1,71 @@
+# agent.py
+import os
+from dotenv import load_dotenv
+from langgraph.graph import StateGraph, MessagesState
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
+
+# Load environment variables (API keys, etc.)
+load_dotenv()
+
+# === TOOLS ===
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers."""
+    return a + b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia and return up to 2 results."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted = "\n\n---\n\n".join(
+        [f"{doc.metadata}\n{doc.page_content[:500]}" for doc in search_docs]
+    )
+    return {"wiki_results": formatted}
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return max 3 results."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted = "\n\n---\n\n".join(
+        [f"{doc.metadata}\n{doc.page_content[:500]}" for doc in search_docs]
+    )
+    return {"web_results": formatted}
+
+tools = [add, wiki_search, web_search]
+
+# === SYSTEM PROMPT ===
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+sys_msg = SystemMessage(content=system_prompt)
+
+# === GRAPH BUILD ===
+def build_graph():
+    # For demo, use OpenAI API if key provided; else default to Gemini
+    from langchain_openai import ChatOpenAI
+    openai_api_key = os.environ.get("OPENAI_API_KEY")
+    if openai_api_key:
+        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=openai_api_key)
+    else:
+        # Gemini fallback (or any other, e.g., groq, etc.)
+        from langchain_google_genai import ChatGoogleGenerativeAI
+        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
+
+    # Bind tools to LLM
+    llm_with_tools = llm.bind_tools(tools)
+
+    # Simple assistant node
+    def assistant(state: MessagesState):
+        # Always include system prompt!
+        messages = [sys_msg] + state["messages"]
+        return {"messages": [llm_with_tools.invoke(messages)]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.set_entry_point("assistant")
+    builder.set_finish_point("assistant")
+
+    return builder.compile()