Use LLM instead of TaskSolver.

lynxkite-lynxscribe/src/lynxkite_lynxscribe/agentic.py

@@ -5,12 +5,12 @@ import os
 import typing
 
 from lynxkite_core import ops
-from lynxscribe.components.task_solver import TaskSolver
 from lynxscribe.components.tool_use import LLM
 from lynxscribe.components.mcp_client import MCPClient
 from lynxscribe.components.rag.rag_chatbot import RAGChatbotResponse
 from lynxscribe.core.llm.base import get_llm_engine
-from lynxscribe.core.models.prompts import Function, Tool
+from lynxscribe.core.models.prompts import Function, Tool, ChatCompletionPrompt
+from lynxscribe.core import router
 
 if typing.TYPE_CHECKING:
     import fastapi
@@ -33,10 +33,7 @@ def gradio_chat(agent: dict):
 
     async def respond(message, chat_history):
         await ag.init()
-        response = await ag.llm.ask(
-            [*chat_history, {"role": "user", "content": message}],
-            expose_tool_calls=True,
-        )
+        response = await ag.ask([*chat_history, {"role": "user", "content": message}])
         async for messages in rag_chatbot_response_to_gradio(response):
             yield messages
 
@@ -56,6 +53,7 @@ def agent_from_dict(agent: dict, default_model: dict) -> "Agent":
         mcp_servers=agent["mcp_servers"],
         agents=[agent_from_dict(a, default_model=model) for a in agent["sub_agents"]],
         model=model,
+        expose_tool_calls=True,
     )
 
 
@@ -153,6 +151,7 @@ class Agent:
         mcp_servers: list[list[str]],
         agents: list["Agent"],
         model: dict,
+        expose_tool_calls: bool,
     ):
         self.name = name
         self.description = description
@@ -160,29 +159,22 @@ class Agent:
         self.mcp_servers = mcp_servers
         self.agents = agents
         self.mcp_client = None
-        self.task_solver = None
+        self.llm = None
         self.llm_engine_params = {**model}
         self.model_name = self.llm_engine_params.pop("model_name")
+        self.expose_tool_calls = expose_tool_calls
 
     async def init(self):
-        if self.task_solver is not None:
+        if self.llm is not None:
             return
         self.mcp_client = MCPClient(*self.mcp_servers)
         await self.mcp_client.aenter()
         agents_as_functions = [agent.as_function() for agent in self.agents]
-        self.task_solver = TaskSolver(
-            llm=get_llm_engine(**self.llm_engine_params),
-            model=self.model_name,
-            initial_messages=[self.prompt],
-            functions=[*self.mcp_client.functions, *agents_as_functions],
-            tool_choice="required",
-            temperature=0.0,
-            max_tool_call_steps=999,
-        )
         self.llm = LLM(
             engine=get_llm_engine(**self.llm_engine_params),
             model=self.model_name,
             tools=[*self.mcp_client.functions, *agents_as_functions],
+            expose_tool_calls=self.expose_tool_calls,
         )
 
     def get_description(self, url: str) -> str:
@@ -206,22 +198,28 @@ class Agent:
 
     async def post(self, request: "fastapi.Request") -> dict:
         if request.state.remaining_path == "chat/completions":
-            request = await request.json()
-            assert not request["stream"]
+            request: dict = await request.json()
+            stream: bool = request.pop("stream", False)
+            request: ChatCompletionPrompt = ChatCompletionPrompt.model_validate(request)
             await self.init()
-            res = await self.task_solver.solve(request["messages"][-1]["content"])
-            return {"choices": [{"message": {"role": "assistant", "content": res}}]}
+            res = await self.ask(request.messages)
+            return await router.rag_chatbot_response_to_http(res, stream)
 
         return {"error": "Not found"}
 
+    async def ask(self, messages: list[str], **kwargs) -> RAGChatbotResponse:
+        await self.init()
+        res = await self.llm.ask([self.prompt, *messages], **kwargs)
+        return res
+
     def as_function(self):
         """A callable that can be used as a tool by another Agent."""
 
-        # Find the value of x given that 4*x^4 = 44. Compute the numerical value.
         async def ask(message):
             print(f"Calling agent {self.name} with message: {message}")
             await self.init()
-            res = await self.task_solver.solve(message)
+            res = await self.ask([message])
+            res = await res.get_answer()
             print(f"Agent {self.name} response: {res}")
             return res
 
@@ -267,7 +265,9 @@ async def rag_chatbot_response_to_gradio(response: RAGChatbotResponse):
                             role="assistant",
                             content="",
                             metadata=dict(
-                                title=f"Using tool: {tool_call.function.name}",
+                                title=f"Asking {tool_call.function.name.removeprefix('ask_').replace('_', ' ')}"
+                                if tool_call.function.name.startswith("ask_")
+                                else f"Using {tool_call.function.name}",
                                 status="pending",
                             ),
                         )