Create ai_agent.py

ai_agent.py

@@ -0,0 +1,68 @@
+from dotenv import load_dotenv
+load_dotenv()
+
+import os
+from langchain_groq import ChatGroq
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langgraph.prebuilt import create_react_agent
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+
+def get_response_from_ai_agent(llm_id, query, allow_search, system_prompt):
+    """
+    Create and invoke an AI agent with optional search capabilities
+    """
+    try:
+        # Initialize LLM with proper configuration
+        llm = ChatGroq(
+            api_key=os.environ.get("GROQ_API_KEY"),
+            model_name=llm_id
+        )
+        
+        # Setup tools based on allow_search flag
+        tools = []
+        if allow_search:
+            tools.append(TavilySearchResults(
+                api_key=os.environ.get("TAVILY_API_KEY"),
+                max_results=2
+            ))
+        
+        # Create the agent
+        agent = create_react_agent(
+            model=llm,
+            tools=tools
+        )
+        
+        # Prepare the initial messages
+        initial_messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=query)
+        ]
+        
+        # Create proper state with messages
+        state = {
+            "messages": initial_messages,
+            "next_steps": [],
+            "structured_response": None
+        }
+        
+        # Invoke agent with proper state
+        response = agent.invoke(state)
+        
+        # Handle response
+        if isinstance(response, dict):
+            # Extract messages from response
+            if "messages" in response:
+                messages = response["messages"]
+                # Get the last AI message
+                ai_messages = [msg for msg in messages if isinstance(msg, AIMessage)]
+                if ai_messages:
+                    return ai_messages[-1].content
+            # Check structured response
+            elif "structured_response" in response:
+                return response["structured_response"]
+        
+        return "I apologize, but I couldn't generate a proper response. Please try again."
+        
+    except Exception as e:
+        print(f"Debug - Error in get_response_from_ai_agent: {str(e)}")
+        raise Exception(f"Agent error: {str(e)}")