File size: 7,739 Bytes
049ae81
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
"""
Chain 03: ReAct Tool Agent
Type: LangChain ReAct Agent
Trigger: action_required | real_world_query | multi-step task
Model: Gemini Pro (best tool-use) / OpenRouter fallback
Tools: Full tool registry (32 tools)
Output: Action results + natural language explanation
"""
from typing import Any, Dict, List, Optional
import time
from langchain.agents import AgentExecutor, create_react_agent
from langchain.prompts import PromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI

from app.config import settings
from app.chains.base import BaseChain, ChainType
from app.tools.registry import ToolRegistry
from app.utils.logging import get_logger

logger = get_logger("react_agent_chain")


class ReActAgentChain(BaseChain):
    """
    ReAct agent chain for multi-step tasks with tool use
    """
    
    chain_type = ChainType.REACT_AGENT
    description = "Multi-step reasoning with tool execution"
    
    REACT_PROMPT = """You are TILLU, a helpful AI assistant. Answer the following question as best you can.

You have access to the following tools:
{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}"""
    
    def __init__(self):
        super().__init__()
        self.llm = None
        self.agent = None
        self._initialize()
    
    def _initialize(self):
        """Initialize LLM and agent"""
        if settings.google_api_key:
            try:
                self.llm = ChatGoogleGenerativeAI(
                    model="gemini-pro",
                    google_api_key=settings.google_api_key,
                    temperature=0.7
                )
            except Exception as e:
                logger.error(f"Failed to initialize Gemini: {e}")
        
        # Fallback to Groq if Gemini not available
        if not self.llm and settings.groq_api_key:
            from langchain_groq import ChatGroq
            self.llm = ChatGroq(
                api_key=settings.groq_api_key,
                model_name="llama-3.1-70b-versatile",
                temperature=0.7
            )
    
    async def execute(
        self,
        input_data: Dict[str, Any],
        context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Execute ReAct agent
        
        Args:
            input_data: Contains 'text' (task/query)
            context: Full context with available tools
            
        Returns:
            Results with tool outputs and explanation
        """
        start_time = time.time()
        
        if not self.llm:
            return {
                "response": {
                    "type": "text",
                    "content": "I'm unable to execute tools right now. Let me answer based on my knowledge.",
                    "structured_data": {}
                },
                "personality_mode": "neutral",
                "chain": self.chain_type.value,
                "model": "fallback",
                "latency_ms": int((time.time() - start_time) * 1000),
                "tokens_used": 0,
                "sources": []
            }
        
        query = input_data.get("text", "")
        
        try:
            # Get tools
            tools = ToolRegistry.get_langchain_tools()
            
            if not tools:
                # No tools available, direct answer
                response = await self.llm.ainvoke([
                    ("human", query)
                ])
                
                return {
                    "response": {
                        "type": "text",
                        "content": response.content,
                        "structured_data": {}
                    },
                    "personality_mode": "sharp",
                    "chain": self.chain_type.value,
                    "model": "groq-70b",
                    "latency_ms": int((time.time() - start_time) * 1000),
                    "tokens_used": 0,
                    "sources": []
                }
            
            # Create ReAct agent
            prompt = PromptTemplate.from_template(self.REACT_PROMPT)
            agent = create_react_agent(self.llm, tools, prompt)
            
            # Execute
            agent_executor = AgentExecutor(
                agent=agent,
                tools=tools,
                verbose=False,
                handle_parsing_errors=True,
                max_iterations=5
            )
            
            result = await agent_executor.ainvoke({"input": query})
            
            elapsed_ms = int((time.time() - start_time) * 1000)
            
            # Extract tool outputs
            intermediate_steps = result.get("intermediate_steps", [])
            tool_outputs = []
            for step in intermediate_steps:
                if len(step) >= 2:
                    action, observation = step[0], step[1]
                    tool_outputs.append({
                        "tool": action.tool if hasattr(action, 'tool') else str(action),
                        "input": action.tool_input if hasattr(action, 'tool_input') else str(action),
                        "output": str(observation)[:500]  # Truncate
                    })
            
            return {
                "response": {
                    "type": "action_result",
                    "content": result.get("output", "Task completed."),
                    "structured_data": {
                        "tools_used": len(tool_outputs),
                        "tool_outputs": tool_outputs,
                        "steps_taken": len(intermediate_steps)
                    }
                },
                "personality_mode": "direct",
                "chain": self.chain_type.value,
                "model": "gemini-pro" if settings.google_api_key else "groq-70b",
                "latency_ms": elapsed_ms,
                "tokens_used": 0,
                "sources": []
            }
            
        except Exception as e:
            logger.error(f"ReAct agent error: {e}")
            
            # Fallback to direct LLM
            try:
                response = await self.llm.ainvoke([("human", query)])
                
                return {
                    "response": {
                        "type": "text",
                        "content": response.content,
                        "structured_data": {"error": str(e), "fallback": True}
                    },
                    "personality_mode": "neutral",
                    "chain": self.chain_type.value,
                    "model": "fallback",
                    "latency_ms": int((time.time() - start_time) * 1000),
                    "tokens_used": 0,
                    "sources": []
                }
            except:
                return {
                    "response": {
                        "type": "text",
                        "content": "I encountered an error while processing your request.",
                        "structured_data": {"error": str(e)}
                    },
                    "personality_mode": "neutral",
                    "chain": self.chain_type.value,
                    "model": "error",
                    "latency_ms": int((time.time() - start_time) * 1000),
                    "tokens_used": 0,
                    "sources": []
                }