Create agent_system.py

agent_system.py

@@ -0,0 +1,147 @@
+import os
+from typing import TypedDict, Annotated, List
+import operator
+
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_huggingface import HuggingFaceEmbeddings, ChatHuggingFace
+from langchain_community.vectorstores import FAISS
+from langchain.tools.retriever import create_retriever_tool
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolExecutor
+
+# --- Configuration ---
+SAVE_PATH = "/data/faiss_index"
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+# Recommended to use a powerful model for agentic tasks
+LLM_REPO_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+
+# --- Agent State Definition ---
+class AgentState(TypedDict):
+    """Defines the state of the agent graph, tracking messages."""
+    messages: Annotated, operator.add]
+
+# --- Knowledge Base and Tools Setup ---
+def create_agent_system():
+    """
+    Initializes the entire agent system, including the knowledge base retriever,
+    specialized tools, and the LangGraph-based agent executor.
+    """
+    print("Initializing Agent System...")
+    
+    # 1. Load the Knowledge Base
+    if not os.path.exists(SAVE_PATH):
+        raise FileNotFoundError(
+            f"FAISS index not found at {SAVE_PATH}. "
+            "Please run knowledge_base.py first to create it."
+        )
+    
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+    vector_store = FAISS.load_local(SAVE_PATH, embeddings, allow_dangerous_deserialization=True)
+    retriever = vector_store.as_retriever(search_kwargs={'k': 3})
+
+    # 2. Create Specialized Retriever Tools for each Agent
+    # The tool descriptions are crucial as they guide the Orchestrator agent.
+    academic_tool = create_retriever_tool(
+        retriever,
+        "academic_retriever",
+        "Searches for information about academics, including curriculum, study materials, timetables, RGPV links, exam papers, and Moodle/e-Library info."
+    )
+    
+    administrative_tool = create_retriever_tool(
+        retriever,
+        "administrative_retriever",
+        "Searches for information about college administration, including fees, scholarships, admissions, rules, regulations, and grievance policies."
+    )
+
+    campus_services_tool = create_retriever_tool(
+        retriever,
+        "campus_services_retriever",
+        "Searches for information about campus services like library hours, bus routes, lab availability, sports facilities, and special academies (Cisco, AWS)."
+    )
+
+    student_life_tool = create_retriever_tool(
+        retriever,
+        "student_life_retriever",
+        "Searches for information about student life, including upcoming events, clubs, cultural festivals, and how to submit complaints or raise issues."
+    )
+
+    tools = [academic_tool, administrative_tool, campus_services_tool, student_life_tool]
+    tool_executor = ToolExecutor(tools)
+
+    # 3. Initialize the LLM
+    # This requires the HUGGINGFACEHUB_API_TOKEN to be set as a secret in the Space.
+    llm = ChatHuggingFace(
+        repo_id=LLM_REPO_ID,
+        task="text-generation",
+        model_kwargs={
+            "max_new_tokens": 1024,
+            "temperature": 0.1,
+            "repetition_penalty": 1.03,
+        },
+    )
+    
+    # Bind the tools to the LLM so it knows how to call them
+    llm_with_tools = llm.bind_tools(tools)
+
+    # 4. Define the LangGraph Nodes
+    def agent_node(state):
+        """The primary node that invokes the LLM to decide the next action."""
+        response = llm_with_tools.invoke(state["messages"])
+        return {"messages": [response]}
+
+    def tool_node(state):
+        """Executes the tool called by the agent and returns the result."""
+        tool_calls = state["messages"][-1].tool_calls
+        tool_messages = tool_executor.batch(tool_calls)
+        return {"messages": tool_messages}
+
+    def should_continue(state):
+        """Conditional edge logic: decides whether to continue or end."""
+        if state["messages"][-1].tool_calls:
+            return "continue"
+        return "end"
+
+    # 5. Build the Graph
+    workflow = StateGraph(AgentState)
+    workflow.add_node("agent", agent_node)
+    workflow.add_node("tools", tool_node)
+    
+    workflow.set_entry_point("agent")
+    
+    workflow.add_conditional_edges(
+        "agent",
+        should_continue,
+        {"continue": "tools", "end": END}
+    )
+    
+    workflow.add_edge("tools", "agent")
+    
+    # 6. Compile the graph into a runnable app
+    agent_executor = workflow.compile()
+    print("Agent System Initialized Successfully.")
+    return agent_executor
+
+def run_query(agent_executor, query, user_info):
+    """
+    Runs a query through the agent system, providing user context.
+    """
+    # Prepend a system message with user context for better personalization
+    system_message = (
+        "You are the 'GGITS Digital Campus Assistant,' a helpful, professional, and reliable AI assistant. "
+        f"You are currently assisting a '{user_info['role']}' named {user_info['email']}. "
+        "Use the available tools to find the most relevant and accurate information from the college's knowledge base. "
+        "If you cannot find an answer, state that the information is not available in your current knowledge base."
+    )
+    
+    messages = [
+        HumanMessage(content=system_message),
+        HumanMessage(content=query)
+    ]
+    
+    # The `stream` method provides real-time output as the agent works
+    final_response = None
+    for chunk in agent_executor.stream({"messages": messages}):
+        if "messages" in chunk:
+            final_response = chunk["messages"][-1]
+
+    return final_response.content if final_response else "I'm sorry, I couldn't process your request."