feat: Finish client

Author: batprem

examples/no-llm-framework/src/no_llm_framework/client/__main__.py

@@ -1,64 +1,43 @@
 import asyncio
 from typing import Literal
-from uuid import uuid4
+
 
 import asyncclick as click
-import httpx
-from a2a.client import A2ACardResolver, A2AClient
-from a2a.types import (
-    Message,
-    MessageSendParams,
-    Part,
-    Role,
-    SendStreamingMessageRequest,
-    SendStreamingMessageSuccessResponse,
-    TaskStatusUpdateEvent,
-    TextPart,
-)
+import colorama
+from no_llm_framework.client.agent import Agent
 
 
 @click.command()
 @click.option('--host', 'host', default='localhost')
 @click.option('--port', 'port', default=9999)
 @click.option('--mode', 'mode', default='streaming')
+@click.option('--question', 'question', required=True)
 async def a_main(
-    host: str, port: int, mode: Literal['completion', 'streaming']
+    host: str,
+    port: int,
+    mode: Literal['completion', 'streaming'],
+    question: str,
 ):
     """Main function to run the A2A Repo Agent client.
 
     Args:
         host (str): The host address to run the server on.
         port (int): The port number to run the server on.
         mode (Literal['completion', 'streaming']): The mode to run the server on.
+        question (str): The question to ask the Agent.
     """  # noqa: E501
-    async with httpx.AsyncClient() as httpx_client:
-        card_resolver = A2ACardResolver(httpx_client, f'http://{host}:{port}')
-        agent_card = await card_resolver.get_agent_card()
-        agent_card.url = f'http://{host}:{port}'
-
-        client = A2AClient(httpx_client, agent_card=agent_card)
-
-        message = MessageSendParams(
-            message=Message(
-                role=Role.user,
-                parts=[Part(TextPart(text='What is Google A2A?'))],
-                messageId=uuid4().hex,
-                taskId=uuid4().hex,
-            )
-        )
-
-        if mode == 'completion':
-            raise NotImplementedError('Completion mode not implemented')
-
-        streaming_request = SendStreamingMessageRequest(params=message)
-        stream_response = client.send_message_streaming(streaming_request)
-        async for chunk in stream_response:
-            if isinstance(
-                chunk.root, SendStreamingMessageSuccessResponse
-            ) and isinstance(chunk.root.result, TaskStatusUpdateEvent):
-                message = chunk.root.result.status.message
-                if message:
-                    print(message.parts[0].root.text, end='', flush=True)
+    agent = Agent(
+        mode='stream',
+        token_stream_callback=None,
+        agent_urls=[f'http://{host}:{port}/'],
+    )
+    async for chunk in agent.stream(question):
+        if chunk.startswith('<Agent name="'):
+            print(colorama.Fore.CYAN + chunk, end='', flush=True)
+        elif chunk.startswith('</Agent>'):
+            print(colorama.Fore.RESET + chunk, end='', flush=True)
+        else:
+            print(chunk, end='', flush=True)
 
 
 def main() -> None:

examples/no-llm-framework/src/no_llm_framework/client/agent.py

@@ -68,6 +68,11 @@ def __init__(
         self.agents_registry: dict[str, AgentCard] = {}
 
     async def get_agents(self) -> tuple[dict[str, AgentCard], str]:
+        """Retrieve agent cards from all agent URLs and render the agent prompt.
+
+        Returns:
+            tuple[dict[str, AgentCard], str]: A dictionary mapping agent names to AgentCard objects, and the rendered agent prompt string.
+        """  # noqa: E501
         async with httpx.AsyncClient() as httpx_client:
             card_resolvers = [
                 A2ACardResolver(httpx_client, url) for url in self.agent_urls
@@ -85,6 +90,14 @@ async def get_agents(self) -> tuple[dict[str, AgentCard], str]:
             return agents_registry, agent_prompt
 
     def call_llm(self, prompt: str) -> str:
+        """Call the LLM with the given prompt and return the response as a string or generator.
+
+        Args:
+            prompt (str): The prompt to send to the LLM.
+
+        Returns:
+            str or Generator[str]: The LLM response as a string or generator, depending on mode.
+        """  # noqa: E501
         if self.mode == 'complete':
             return stream_llm(prompt)
 
@@ -94,10 +107,31 @@ def call_llm(self, prompt: str) -> str:
         return result
 
     async def decide(
-        self, question: str, agents_prompt: str
+        self,
+        question: str,
+        agents_prompt: str,
+        called_agents: list[dict] | None = None,
     ) -> Generator[str, None]:
+        """Decide which agent(s) to use to answer the question.
+
+        Args:
+            question (str): The question to answer.
+            agents_prompt (str): The prompt describing available agents.
+            called_agents (list[dict] | None): Previously called agents and their answers.
+
+        Returns:
+            Generator[str, None]: The LLM's response as a generator of strings.
+        """  # noqa: E501
+        if called_agents:
+            call_agent_prompt = agent_answer_template.render(
+                called_agents=called_agents
+            )
+        else:
+            call_agent_prompt = ''
         prompt = decide_template.render(
-            question=question, agent_prompt=agents_prompt
+            question=question,
+            agent_prompt=agents_prompt,
+            call_agent_prompt=call_agent_prompt,
         )
         return self.call_llm(prompt)
 
@@ -116,6 +150,15 @@ def extract_agents(self, response: str) -> list[dict]:
     async def send_message_to_an_agent(
         self, agent_card: AgentCard, message: str
     ):
+        """Send a message to a specific agent and yield the streaming response.
+
+        Args:
+            agent_card (AgentCard): The agent to send the message to.
+            message (str): The message to send.
+
+        Yields:
+            str: The streaming response from the agent.
+        """
         async with httpx.AsyncClient() as httpx_client:
             client = A2AClient(httpx_client, agent_card=agent_card)
             message = MessageSendParams(
@@ -137,65 +180,73 @@ async def send_message_to_an_agent(
                         yield message.parts[0].root.text
 
     async def stream(self, question: str):
-        agents_registry, agent_prompt = await self.get_agents()
-        response = ''
-        for chunk in await self.decide(question, agent_prompt):
-            response += chunk
-            if self.token_stream_callback:
-                self.token_stream_callback(chunk)
-            yield chunk
-
-        agents = self.extract_agents(response)
+        """Stream the process of answering a question, possibly involving multiple agents.
+
+        Args:
+            question (str): The question to answer.
+
+        Yields:
+            str: Streaming output, including agent responses and intermediate steps.
+        """  # noqa: E501
         agent_answers: list[dict] = []
-        for agent in agents:
-            agent_response = ''
-            agent_card = agents_registry[agent['name']]
-            yield f'<Agent name="{agent["name"]}">\n'
-            async for chunk in self.send_message_to_an_agent(
-                agent_card, agent['prompt']
+        for _ in range(3):
+            agents_registry, agent_prompt = await self.get_agents()
+            response = ''
+            for chunk in await self.decide(
+                question, agent_prompt, agent_answers
             ):
-                agent_response += chunk
+                response += chunk
                 if self.token_stream_callback:
                     self.token_stream_callback(chunk)
                 yield chunk
-            yield '</Agent>\n'
-            match = re.search(
-                r'<Answer>(.*?)</Answer>', agent_response, re.DOTALL
-            )
-            answer = match.group(1).strip() if match else agent_response
-            if answer:
-                agent_answers.append(
-                    {
-                        'name': agent['name'],
-                        'prompt': agent['prompt'],
-                        'answer': answer,
-                    }
-                )
+
+            agents = self.extract_agents(response)
+            if agents:
+                for agent in agents:
+                    agent_response = ''
+                    agent_card = agents_registry[agent['name']]
+                    yield f'<Agent name="{agent["name"]}">\n'
+                    async for chunk in self.send_message_to_an_agent(
+                        agent_card, agent['prompt']
+                    ):
+                        agent_response += chunk
+                        if self.token_stream_callback:
+                            self.token_stream_callback(chunk)
+                        yield chunk
+                    yield '</Agent>\n'
+                    match = re.search(
+                        r'<Answer>(.*?)</Answer>', agent_response, re.DOTALL
+                    )
+                    answer = match.group(1).strip() if match else agent_response
+                    agent_answers.append(
+                        {
+                            'name': agent['name'],
+                            'prompt': agent['prompt'],
+                            'answer': answer,
+                        }
+                    )
             else:
-                print('<Answer> tag not found')
-        print(agent_answers)
+                return
 
 
 if __name__ == '__main__':
     import asyncio
     import colorama
 
     async def main():
+        """Main function to run the A2A Repo Agent client."""
         agent = Agent(
             mode='stream',
             token_stream_callback=None,
             agent_urls=['http://localhost:9999/'],
         )
-        agents_registry, agent_prompt = await agent.get_agents()
-        # agent_card = agents_registry['A2A Protocol Agent']
+
         async for chunk in agent.stream('What is A2A protocol?'):
             if chunk.startswith('<Agent name="'):
                 print(colorama.Fore.CYAN + chunk, end='', flush=True)
             elif chunk.startswith('</Agent>'):
                 print(colorama.Fore.RESET + chunk, end='', flush=True)
             else:
                 print(chunk, end='', flush=True)
-        # async for chunk in stream_response:
-        #     print(chunk, end='', flush=True)
 
     asyncio.run(main())

examples/no-llm-framework/src/no_llm_framework/client/agent_answer.jinja

@@ -0,0 +1,5 @@
+Previous agents have been called. {% for agent in called_agents %}
+- Agent: {{ agent.name }}
+- Prompt: {{ agent.prompt }}
+- Answer: {{ agent.answer }}
+{% endfor %}

examples/no-llm-framework/src/no_llm_framework/client/decide.jinja

@@ -3,7 +3,7 @@ The question is:
 
 {{ question }}
 
-{{ called_agents }}
+{{ call_agent_prompt }}
 
 {{ agent_prompt }}
 

examples/no-llm-framework/src/no_llm_framework/server/agent.py

@@ -53,10 +53,18 @@ def __init__(
         self.mcp_url = mcp_url
 
     def call_llm(self, prompt: str) -> Generator[str, None]:
+        """Call the LLM with the given prompt and return a generator of responses.
+
+        Args:
+            prompt (str): The prompt to send to the LLM.
+
+        Returns:
+            Generator[str, None]: A generator yielding the LLM's response.
+        """   # noqa: E501
         return stream_llm(prompt)
 
     async def decide(
-        self, question: str, called_tools: list[dict]
+        self, question: str, called_tools: list[dict] | None = None
     ) -> Generator[str, None]:
         """Decide which tool to use to answer the question.
 
@@ -107,6 +115,14 @@ async def call_tool(self, tools: list[dict]) -> list[CallToolResult]:
         )
 
     async def stream(self, question: str) -> AsyncGenerator[str]:
+        """Stream the process of answering a question, possibly involving tool calls.
+
+        Args:
+            question (str): The question to answer.
+
+        Yields:
+            dict: Streaming output, including intermediate steps and final result.
+        """  # noqa: E501
         called_tools = []
         for i in range(10):
             yield {