|
1 | | -from typing import Dict, Generator |
| 1 | +import uuid |
| 2 | +from abc import ABC, abstractmethod |
| 3 | +from typing import Dict, Generator, List, Optional |
2 | 4 |
|
3 | 5 | from application.agents.llm_handler import get_llm_handler |
4 | 6 | from application.agents.tools.tool_action_parser import ToolActionParser |
5 | 7 | from application.agents.tools.tool_manager import ToolManager |
6 | 8 |
|
7 | 9 | from application.core.mongo_db import MongoDB |
8 | 10 | from application.llm.llm_creator import LLMCreator |
| 11 | +from application.logging import build_stack_data, log_activity, LogContext |
| 12 | +from application.retriever.base import BaseRetriever |
9 | 13 |
|
10 | 14 |
|
11 | | -class BaseAgent: |
| 15 | +class BaseAgent(ABC): |
12 | 16 | def __init__( |
13 | 17 | self, |
14 | | - endpoint, |
15 | | - llm_name, |
16 | | - gpt_model, |
17 | | - api_key, |
18 | | - user_api_key=None, |
19 | | - decoded_token=None, |
20 | | - attachments=None, |
| 18 | + endpoint: str, |
| 19 | + llm_name: str, |
| 20 | + gpt_model: str, |
| 21 | + api_key: str, |
| 22 | + user_api_key: Optional[str] = None, |
| 23 | + prompt: str = "", |
| 24 | + chat_history: Optional[List[Dict]] = None, |
| 25 | + decoded_token: Optional[Dict] = None, |
| 26 | + attachments: Optional[str]=None, |
21 | 27 | ): |
22 | 28 | self.endpoint = endpoint |
| 29 | + self.llm_name = llm_name |
| 30 | + self.gpt_model = gpt_model |
| 31 | + self.api_key = api_key |
| 32 | + self.user_api_key = user_api_key |
| 33 | + self.prompt = prompt |
| 34 | + self.decoded_token = decoded_token or {} |
| 35 | + self.user: str = decoded_token.get("sub") |
| 36 | + self.tool_config: Dict = {} |
| 37 | + self.tools: List[Dict] = [] |
| 38 | + self.tool_calls: List[Dict] = [] |
| 39 | + self.chat_history: List[Dict] = chat_history if chat_history is not None else [] |
23 | 40 | self.llm = LLMCreator.create_llm( |
24 | 41 | llm_name, |
25 | 42 | api_key=api_key, |
26 | 43 | user_api_key=user_api_key, |
27 | 44 | decoded_token=decoded_token, |
28 | 45 | ) |
29 | 46 | self.llm_handler = get_llm_handler(llm_name) |
30 | | - self.gpt_model = gpt_model |
31 | | - self.tools = [] |
32 | | - self.tool_config = {} |
33 | | - self.tool_calls = [] |
34 | | - self.attachments = attachments or [] |
| 47 | + set.attachments = attachments or [] |
35 | 48 |
|
36 | | - def gen(self, *args, **kwargs) -> Generator[Dict, None, None]: |
37 | | - raise NotImplementedError('Method "gen" must be implemented in the child class') |
| 49 | + @log_activity() |
| 50 | + def gen( |
| 51 | + self, query: str, retriever: BaseRetriever, log_context: LogContext = None |
| 52 | + ) -> Generator[Dict, None, None]: |
| 53 | + yield from self._gen_inner(query, retriever, log_context) |
| 54 | + |
| 55 | + @abstractmethod |
| 56 | + def _gen_inner( |
| 57 | + self, query: str, retriever: BaseRetriever, log_context: LogContext |
| 58 | + ) -> Generator[Dict, None, None]: |
| 59 | + pass |
38 | 60 |
|
39 | 61 | def _get_user_tools(self, user="local"): |
40 | 62 | mongo = MongoDB.get_client() |
@@ -111,14 +133,12 @@ def _execute_tool_action(self, tools_dict, call): |
111 | 133 | for param, details in action_data[param_type]["properties"].items(): |
112 | 134 | if param not in call_args and "value" in details: |
113 | 135 | target_dict[param] = details["value"] |
114 | | - |
115 | 136 | for param, value in call_args.items(): |
116 | 137 | for param_type, target_dict in param_types.items(): |
117 | 138 | if param_type in action_data and param in action_data[param_type].get( |
118 | 139 | "properties", {} |
119 | 140 | ): |
120 | 141 | target_dict[param] = value |
121 | | - |
122 | 142 | tm = ToolManager(config={}) |
123 | 143 | tool = tm.load_tool( |
124 | 144 | tool_data["name"], |
@@ -153,3 +173,79 @@ def _execute_tool_action(self, tools_dict, call): |
153 | 173 | self.tool_calls.append(tool_call_data) |
154 | 174 |
|
155 | 175 | return result, call_id |
| 176 | + |
| 177 | + def _build_messages( |
| 178 | + self, |
| 179 | + system_prompt: str, |
| 180 | + query: str, |
| 181 | + retrieved_data: List[Dict], |
| 182 | + ) -> List[Dict]: |
| 183 | + docs_together = "\n".join([doc["text"] for doc in retrieved_data]) |
| 184 | + p_chat_combine = system_prompt.replace("{summaries}", docs_together) |
| 185 | + messages_combine = [{"role": "system", "content": p_chat_combine}] |
| 186 | + |
| 187 | + for i in self.chat_history: |
| 188 | + if "prompt" in i and "response" in i: |
| 189 | + messages_combine.append({"role": "user", "content": i["prompt"]}) |
| 190 | + messages_combine.append({"role": "assistant", "content": i["response"]}) |
| 191 | + if "tool_calls" in i: |
| 192 | + for tool_call in i["tool_calls"]: |
| 193 | + call_id = tool_call.get("call_id") or str(uuid.uuid4()) |
| 194 | + |
| 195 | + function_call_dict = { |
| 196 | + "function_call": { |
| 197 | + "name": tool_call.get("action_name"), |
| 198 | + "args": tool_call.get("arguments"), |
| 199 | + "call_id": call_id, |
| 200 | + } |
| 201 | + } |
| 202 | + function_response_dict = { |
| 203 | + "function_response": { |
| 204 | + "name": tool_call.get("action_name"), |
| 205 | + "response": {"result": tool_call.get("result")}, |
| 206 | + "call_id": call_id, |
| 207 | + } |
| 208 | + } |
| 209 | + |
| 210 | + messages_combine.append( |
| 211 | + {"role": "assistant", "content": [function_call_dict]} |
| 212 | + ) |
| 213 | + messages_combine.append( |
| 214 | + {"role": "tool", "content": [function_response_dict]} |
| 215 | + ) |
| 216 | + messages_combine.append({"role": "user", "content": query}) |
| 217 | + return messages_combine |
| 218 | + |
| 219 | + def _retriever_search( |
| 220 | + self, |
| 221 | + retriever: BaseRetriever, |
| 222 | + query: str, |
| 223 | + log_context: Optional[LogContext] = None, |
| 224 | + ) -> List[Dict]: |
| 225 | + retrieved_data = retriever.search(query) |
| 226 | + if log_context: |
| 227 | + data = build_stack_data(retriever, exclude_attributes=["llm"]) |
| 228 | + log_context.stacks.append({"component": "retriever", "data": data}) |
| 229 | + return retrieved_data |
| 230 | + |
| 231 | + def _llm_gen(self, messages: List[Dict], log_context: Optional[LogContext] = None): |
| 232 | + resp = self.llm.gen_stream( |
| 233 | + model=self.gpt_model, messages=messages, tools=self.tools |
| 234 | + ) |
| 235 | + if log_context: |
| 236 | + data = build_stack_data(self.llm) |
| 237 | + log_context.stacks.append({"component": "llm", "data": data}) |
| 238 | + return resp |
| 239 | + |
| 240 | + def _llm_handler( |
| 241 | + self, |
| 242 | + resp, |
| 243 | + tools_dict: Dict, |
| 244 | + messages: List[Dict], |
| 245 | + log_context: Optional[LogContext] = None, |
| 246 | + ): |
| 247 | + resp = self.llm_handler.handle_response(self, resp, tools_dict, messages) |
| 248 | + if log_context: |
| 249 | + data = build_stack_data(self.llm_handler) |
| 250 | + log_context.stacks.append({"component": "llm_handler", "data": data}) |
| 251 | + return resp |
0 commit comments