cognitrix 0.2.4__tar.gz → 0.2.5__tar.gz

{cognitrix-0.2.4 → cognitrix-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognitrix
-Version: 0.2.4
+Version: 0.2.5
 Summary: Package for creating AI Agents using llms
 Home-page: https://github.com/theonlyamos/cognitrix
 License: Apache2
@@ -28,10 +28,10 @@ Requires-Dist: aiohttp (>=3.9.3,<4.0.0)
 Requires-Dist: anthropic (>=0.17.0,<0.18.0)
 Requires-Dist: beautifulsoup4 (>=4.12.3,<5.0.0)
 Requires-Dist: clarifai (>=10.1.0,<11.0.0)
-Requires-Dist: cohere (>=4.50,<5.0)
+Requires-Dist: cohere (>=5.5.6,<6.0.0)
 Requires-Dist: deepgram-sdk (==3.*)
+Requires-Dist: flask (>=3.0.3,<4.0.0)
 Requires-Dist: google-generativeai (>=0.4.0,<0.5.0)
-Requires-Dist: google-search-results (>=2.4.2,<3.0.0)
 Requires-Dist: groq (>=0.4.1,<0.5.0)
 Requires-Dist: ollama (>=0.2.0,<0.3.0)
 Requires-Dist: openai (>=1.12.0,<2.0.0)
@@ -42,9 +42,10 @@ Requires-Dist: pydantic (>=2.6.2,<3.0.0)
 Requires-Dist: pydub (>=0.25.1,<0.26.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
-Requires-Dist: serpapi (>=0.1.5,<0.2.0)
+Requires-Dist: tavily-python (>=0.3.3,<0.4.0)
 Requires-Dist: together (>=0.2.11,<0.3.0)
 Requires-Dist: uuid (>=1.30,<2.0)
+Requires-Dist: uvicorn (>=0.30.1,<0.31.0)
 Requires-Dist: xmltodict (>=0.13.0,<0.14.0)
 Project-URL: Bug Tracker, https://github.com/theonlyamos/cognitrix/issues
 Project-URL: Repository, https://github.com/theonlyamos/cognitrix

{cognitrix-0.2.4 → cognitrix-0.2.5}/cognitrix/agents/base.py

@@ -1,19 +1,17 @@
 import inspect
 import json
-import sys
 import uuid
 import asyncio
 import logging
 import aiofiles
-from pathlib import Path
 from threading import Thread
-from typing import List, Optional, Self, TypeAlias, Union, Type, Any
-from dataclasses import dataclass
+from datetime import datetime
+from typing import Dict, List, Optional, Self, TypeAlias, Union, Type, Any
 
 from pydantic import BaseModel, Field
 
 from cognitrix.tasks import Task
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from cognitrix.tools.base import Tool
 from cognitrix.utils import extract_json, json_return_format
 from cognitrix.agents.templates import AUTONOMOUSE_AGENT_2
@@ -27,43 +25,112 @@ AgentList: TypeAlias = List['Agent']
 
 class Agent(BaseModel):
     name: str = Field(default='Agent')
+    """Name of the agent"""
+    
     llm: LLM
+    """LLM Provider to use for the agent"""
+    
     tools: List[Tool] = Field(default_factory=list)
+    """List of tools to be use by the agent"""
+    
     prompt_template: str = Field(default=AUTONOMOUSE_AGENT_2)
+    """Agent's prompt template"""
+    
     verbose: bool = Field(default=False)
+    """Set agent verbosity"""
+    
     sub_agents: AgentList = Field(default_factory=list)
+    """List of sub agents which can be called by this agent"""
+    
     task: Optional[Task] = None
+    """The task to be completed by the agent"""
+    
     is_sub_agent: bool = Field(default=False)
+    """Whether this agent is a sub agent for another agent"""
+    
     id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    """Unique id for the agent"""
+    
     parent_id: Optional[str] = None
+    """Id of this agent's parent agent (if it's a sub agent)"""
+    
     autostart: bool = False
+    """Whether the agent should start running as soon as it's created"""
 
     @property
     def available_tools(self) -> List[str]:
         return [tool.name for tool in self.tools]
+    
+    def format_tools_for_llm(self):
+        TYPE_HINTS = {
+            "<class 'int'>": 'integer',
+            "<class 'str'>": 'string',
+            "<class 'float'>": 'float',
+            "<class 'dict'>": 'object',
+            "<class 'list'>": 'array',
+            "typing.Dict": 'object',
+            "typing.List": 'array',
+            "typing.Optional[int]": 'integer',
+            "typing.Optional[str]": 'integer',
+            "typing.Optional[float]": 'float',
+        }
+        llm_tools: list[dict[str, Any]] = []
+        classes = []
+        for tool in self.tools:
+            tool_parameters = tool.parameters
+            if not tool_parameters:
+                func_signatures = inspect.signature(tool.run)
+                tool_parameters = func_signatures.parameters
+            
+            parameters = tool_parameters.keys()
+            tool_details = {
+                'tool': tool,
+                'name': tool.name,
+                'description': tool.description,
+                'parameters': {},
+                'required': []
+            }
+
+            if not 'args' in parameters:
+                tool_details['parameters'] = {name: TYPE_HINTS.get(str(param.annotation), 'string') for name, param in tool_parameters.items()}
+                tool_details['required'] = [name for name, param in tool_parameters.items() if param.default is inspect._empty]
+            
+            llm_tools.append(tool_details)
+
+        self.llm.format_tools(llm_tools)
 
     def format_system_prompt(self):
         tools_str = self._format_tools_string()
         subagents_str = self._format_subagents_string()
         llms_str = self._format_llms_string()
 
-        prompt = self.prompt_template
+        today = (datetime.now()).strftime("%a %b %d %Y") 
+        prompt = f"Today is {today}.\n\n"
+        prompt += self.prompt_template
         prompt = prompt.replace("{name}", self.name)
+        
+        # if not self.llm.supports_tool_use:
         prompt = prompt.replace("{tools}", tools_str)
         prompt = prompt.replace("{subagents}", subagents_str)
+        
+        # if not self.llm.supports_tool_use:
         prompt = prompt.replace("{available_tools}", json.dumps(self.available_tools))
         prompt = prompt.replace("{llms}", llms_str)
         prompt = prompt.replace("{return_format}", json_return_format)
 
         if 'json' not in prompt.lower():
             prompt += f"\n{json_return_format}"
-        
+
         self.llm.system_prompt = prompt
+        
+        self.format_tools_for_llm()
 
     def _format_tools_string(self) -> str:
         return "You have access to the following Tools:\n" + "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
 
     def _format_subagents_string(self) -> str:
+        if not len(self.sub_agents):
+            return ''
         subagents_str = "Available Subagents:\n"
         subagents_str += "\n".join([f"-- {agent.name}: {agent.task.description}" for agent in self.sub_agents if agent.task])
         subagents_str += "\nYou should always use a subagent for a task if there is one specifically created for that task."
@@ -83,6 +150,7 @@ class Agent(BaseModel):
             if self.is_sub_agent:
                 print("=======is sub agent===========")
                 print(processed_query)
+
             result = processed_query['result']
             if isinstance(result, list):
                 if result[0] == 'image':
@@ -115,54 +183,44 @@ class Agent(BaseModel):
     def get_tool_by_name(self, name: str) -> Optional[Tool]:
         return next((tool for tool in self.tools if tool.name.lower() == name.lower()), None)
 
-    async def process_response(self, response: str|dict) -> Union[dict, str]:
-        # response = response.replace("'", '"')
-        response_data = response
-        if isinstance(response, str):
-            response = response.replace('\\n', '')
-            response = response.replace("'", "\"")
-            # response = response.replace('"', '\\"')
-            response_data = extract_json(response)
-        
-
+    async def call_tools(self, tool_calls: list) -> Union[dict, str]:
+ 
         try:
-            if isinstance(response_data, dict):
-                # final_result_keys = ['final_answer', 'function_call_result', 'respons']
-
-                if response_data['type'].replace('\\', '') != 'function_call':
-                    return response_data['result']
-
-                tool = self.get_tool_by_name(response_data['function'])
-                if isinstance(response_data['arguments'], dict):
-                    response_data['arguments'] = list(response_data['arguments'].values())
-
-                if response_data['function'].lower() == 'create agents':
-                    response_data['arguments'] = [*response_data['arguments'], self.id]
-
-                if not tool:
-                    raise Exception(f"Tool {response_data['function']} not found")
-
-                print(f"\nRunning tool '{tool.name.title()}' with parameters: {response_data['arguments']}")
-
-                if 'sub agent' in tool.name.lower():
-                    response_data['arguments'].append(self)
+            if tool_calls:
+                tool_calls_result = []
+                for t in tool_calls:
+                    tool = self.get_tool_by_name(t['name'])
                     
-                if tool.name.lower() == 'create sub agent':
-                    result = await tool.arun(*response_data['arguments'])
-                else:
-                    result = tool.run(*response_data['arguments'])
-
-                response_json = {
-                    'type': 'function_call_result',
-                    'result': result
+                    if not tool:
+                        print(f"Tool '{t['name']}' not found")
+                        raise Exception(f"Tool '{t['name']}' not found")
+                    
+                    print(f"\nRunning tool '{tool.name.title()}' with parameters: {t['arguments']}")
+
+                    if 'sub agent' in tool.name.lower():
+                        t['arguments']['parent'] = self
+                        
+                    if tool.name.lower() == 'create sub agent':
+                        result = await tool.arun(*t['arguments'])
+                    else:
+                        result = tool.run(**t['arguments'])
+
+                    tool_calls_result.append({
+                        'tool_call_id': t['name'],
+                        'role': 'tool',
+                        'name': tool.name,
+                        'content': result
+                    })
+                
+                return {
+                    'type': 'tool_calls_result',
+                    'result': f"The results of the tool_calls are {json.dumps(tool_calls_result)}"
                 }
-
-                return response_json
             else:
                 raise Exception('Not a json object')
         except Exception as e:
-            # logger.exception(e)
-            return response_data
+            logger.exception(e)
+            return str(e)
 
     def add_tool(self, tool: Tool):
         self.tools.append(tool)
@@ -208,18 +266,22 @@ class Agent(BaseModel):
                 response: Any = self.llm(full_prompt)
                 
                 self.llm.chat_history.append(full_prompt)
-                self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response})
-
-                if self.verbose:
-                    print(response)
-
-                result: dict[Any, Any] | str = asyncio.run(self.process_response(response))
+                
+                if response.text:
+                    self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response.text})
+                    print(f"\n{self.name}: {response.text}")
+                
+                if response.tool_calls:
+                    result: dict[Any, Any] | str = asyncio.run(self.call_tools(response.tool_calls))
 
-                if isinstance(result, dict) and result['type'] == 'function_call_result':
-                    query = result
+                    if isinstance(result, dict) and result['type'] == 'tool_calls_result':
+                        query = result
+                    else:
+                        print(result)
                 else:
-                    print(f"\n{self.name}: {result}")
                     query = input("\nUser (q to quit): ")
+
+                # query = input("\nUser (q to quit): ")
                 
                 self.save_session(session)
 
@@ -248,9 +310,9 @@ class Agent(BaseModel):
             if self.verbose:
                 print(response)
 
-            processsed_response: dict[Any, Any] | str = asyncio.run(self.process_response(response))
+            processsed_response: dict[Any, Any] | str = asyncio.run(self.call_tools(response))
 
-            if isinstance(processsed_response, dict) and processsed_response['type'] == 'function_call_result':
+            if isinstance(processsed_response, dict) and processsed_response['type'] == 'tool_calls_result':
                 query = processsed_response
             else:
                 if isinstance(processsed_response, str):
@@ -287,8 +349,8 @@ class Agent(BaseModel):
             if parent.verbose:
                 print(response)
                 
-            agent_result = asyncio.run(self.process_response(response))
-            if isinstance(agent_result, dict) and agent_result['type'] == 'function_call_result':
+            agent_result = asyncio.run(self.call_tools(response))
+            if isinstance(agent_result, dict) and agent_result['type'] == 'tool_calls_result':
                 query = agent_result
             else:
                 print(f"\n--{self.name}: {agent_result}")
@@ -297,7 +359,7 @@ class Agent(BaseModel):
                 parent_response: Any = parent.llm(parent_prompt)
                 parent.llm.chat_history.append(parent_prompt)
                 parent.llm.chat_history.append({'role': 'assistant', 'type': 'text', 'message': parent_response})
-                parent_result = asyncio.run(parent.process_response(parent_response))
+                parent_result = asyncio.run(parent.call_tools(parent_response))
                 print(f"\n\n{parent.name}: {parent_result}")
                 query = ""
 

{cognitrix-0.2.4 → cognitrix-0.2.5}/cognitrix/agents/templates.py

@@ -74,9 +74,9 @@ Context: The AI assistant has access to the user's computer and the internet. Th
 The AI assistant will answer the user's question to the best of its ability, using its knowledge and access to tools provided by the user.
 The AI assistant has access to the following tools.
 
-Remember, functions calls will be processed by the user and the result returned to the AI Asssistant as the next input.
-A function name is a name of a tool available to the AI Assistant.
-Whenever there is a function call, always wait for the answer from the user. Do not try to answer that query yourself.
+Remember, tools calls will be processed by the user and the result returned to the AI Asssistant as the next input.
+A tool name is a name of a tool available to the AI Assistant.
+Whenever there is a Tool Call, always wait for the answer from the user. Do not try to answer that query yourself.
 Only call tools available to the AI Assistant.
 
 Tools:
@@ -88,15 +88,15 @@ User: Answer my question: What is the capital of France?
 AI Assistant: {"type": "final_answer", "result": "Paris"}
 
 User: What is the capital of the country with the highest population in the world?
-AI Assistant: {"type": "function_call", "function": "Internet Search", "arguments": ["current country with highest population in the world"]}
+AI Assistant: {"type": "tool_calls", "tool_calls": [{"name": "Internet Search", "arguments": {"query": "current country with highest population in the world"}}]}
 
-User: {"type": "function_call_result", "result": "China, Beijing"}
+User: {"type": "tool_calls_result", "results": [{"name": "Internet Search", "result": "China, Beijing"}]}
 AI Assistant: {"type": "final_answer", "result": "The current highest country with highest population in the world is China, Beijing"}
 
 User: Take a screenshot of my screen.
-AI Assistant: {"type": "function_call", "function": "Screenshot", "arguments": []}
+AI Assistant: {"type": "tool_calls", "tool_calls": [{"name": "Screenshot", "arguments": {}}]}
 
-User: {"type": "function_call_result", "result": "C:/Users/my_username/Desktop/screenshot.png"}
+User: {"type": "tool_calls_result", "results": [{"name": "Screenshot", "result": "C:/Users/my_username/Desktop/screenshot.png"}]}
 AI Assistant: {"type": "final_answer", "result": "Screenshot saved to C:/Users/my_username/Desktop/screenshot.png"}
 
 The AI Assistant's output should always be in a json format.
@@ -106,16 +106,15 @@ json.loads()
 Return the response in the following format only:
 {
   "type": "final_answer",
-  "result": "
+  "result": "The is the final answer"
 }
-if it's the final anwers or
+if it's the final answer or
 {
-  "type": "function_call",
-  "function": "",
-  "arguments": []
+  "type": "tool_calls",
+  "tool_calls": [{"name": "<tool_name>", "arguments": {"<arg_name>": "<arg_value>", ...}}, ...]
 }
 if the assistant needs to use a tool to answer the user's query.
-Don't forget to present the answer to a function call to the user in an informative manner.
+Don't forget to present the answer to a Tool Call to the user in an informative manner.
 You should always break down complex tasks into smaller easier ones and perform them one by one. 
 You should always check if a task is complete by taking a screenshot of the screen.
 Begin:
@@ -177,7 +176,7 @@ Process:
    - Iterate until the user's requested task is completed successfully.
    - Provide feedback or status updates to the user to indicate progress, errors, or completion of the task.
 
-Remember, functions calls will be processed by the user and the result returned to the AI agent as the next input. A function name is a name of a tool available to the AI agent. Whenever there is a function call, always wait for the answer from the user. Do not try to complete that task yourself. Only call tools available to the AI agent.
+Remember, tool calls will be processed by the user and the result returned to the AI agent as the next input. A tool name is a name of a tool available to the AI agent. Whenever there is a Tool Call, always wait for the answer from the user. Do not try to complete that task yourself. Only call tools available to the AI agent.
 
 Constraints:
 - Operate within the boundaries of the computer screen and the available UI elements.
@@ -197,14 +196,13 @@ Return the response in the following format only:
 if it's the final result of the task, or
 
 {
-  "type": "function_call",
-  "function": "<Function Name>",
-  "arguments": ["<arg1>", "<arg2>", ...]
+  "type": "tool_calls",
+  "tool_calls": [{"name": "<tool_name>", "arguments": {"<arg_name>": "<arg_value>", ...}}, ...]
 }
 
 if the agent needs to use a tool to complete the task.
 
-Don't forget to present the result of a function call to the user in an informative manner.
+Don't forget to present the result of a Tool Call to the user in an informative manner.
 
 Remember, as an autonomous AI agent, your goal is to efficiently navigate the computer's user interface, interact with UI elements, 
 and utilize the provided tools to accomplish the tasks requested by the user. 
@@ -235,21 +233,18 @@ You are an AI agent designed to assist users with various tasks on their compute
 5. **Communicate:** 
     * **Format your response as JSON:** Structure your response in one of the following valid JSON formats:
         * **Final result:**
-        ```json
         {
             "type": "final_answer",
             "result": "<result>"
         }
-        ```
-        * **Function call:**
-        ```json
+        
+        * **Tool call:**
         {
-            "type": "function_call",
-            "function": "<Function Name>",
-            "arguments": ["<arg1>", "<arg2>", ...]
+            "type": "tool_calls",
+            "tool_calls": [{"name": "<tool_name>", "arguments": {"<arg_name>": "<arg_value>", ...}}, ...]
         }
-        ```
-    * **Keep the user informed:** Provide clear and informative communication to the user about your progress and any actions you take.
+
+* **Keep the user informed:** Provide clear and informative communication to the user about your progress and any actions you take.
 
 **Additional Considerations:**
 
@@ -313,13 +308,12 @@ Example 1 (Final Answer with Chain of Thought):
     "result": "The capital of France is Paris."
 }
 
-Example 2 (Function Call with Chain of Thought):
+Example 2 (Tool Call with Chain of Thought):
 {
     "observation": "The user asked me to search for information about artificial intelligence on Wikipedia.",
     "thought": "Step 1) To search for information on Wikipedia, I need to use the 'search_wikipedia' tool.\\nStep 2) The relevant argument for this tool is the search query, which in this case is 'artificial intelligence'.",
-    "type": "function_call",
-    "function": "search_wikipedia",
-    "arguments": ["artificial intelligence"]
+    "type": "tool_calls",
+    "tool_calls": [{"name": "Search Wikipedia", "arguments": {"query": "artificial intelligence"}}]
 }
 
 Example 3 (Final Answer with Chain of Thought):
@@ -330,13 +324,12 @@ Example 3 (Final Answer with Chain of Thought):
     "result": "The area of the rectangle is 15 square meters."
 }
 
-Example 4 (Function Call with Chain of Thought):
+Example 4 (Tool Call with Chain of Thought):
 {
     "observation": "The user asked me to find the current weather forecast for New York City.",
-    "thought": "Step 1) To find the weather forecast for a specific location, I need to use the 'get_weather_forecast' tool.\\nStep 2) The relevant argument for this tool is the location, which in this case is 'New York City'.",
-    "type": "function_call",
-    "function": "get_weather_forecast",
-    "arguments": ["New York City"]
+    "thought": "Step 1) To find the weather forecast for a specific location, I need to use the 'internet_search' tool.\\nStep 2) The relevant argument for this tool is the location, which in this case is 'New York City'.",
+    "type": "tool_calls",
+    "tool_calls": [{"name": "Internet Search", "arguments": {"query": "weather forecast for New York City today"}}]
 }
 
 Example 5 (Final Answer with Chain of Thought):
@@ -347,7 +340,7 @@ Example 5 (Final Answer with Chain of Thought):
     "result": "25 degrees Celsius is equivalent to 77 degrees Fahrenheit."
 }
 
-The only keys allowed in the json response are ["observervation", "tought", "type", "function", "arguments", "result"].
+The only keys allowed in the json response are ["observation", "thought", "type", "tool_calls", "result"].
 """
 
 AUTONOMOUSE_AGENT_2_JSON_REMINDER = """

{cognitrix-0.2.4 → cognitrix-0.2.5}/cognitrix/cli.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from argparse import Namespace
 
 from cognitrix.llms import (
-    Cohere, OpenAI,LLM
+    Cohere, Clarifai, LLM
 )
 
 from cognitrix.agents import AIAssistant, Agent
@@ -128,9 +128,6 @@ def start(args: Namespace):
         elif args.agents:
             list_agents()              #type: ignore
             sys.exit()
-        elif args.tools:
-            list_tools()
-            sys.exit()
         elif args.sessions:
             list_sessions()
             sys.exit()
@@ -138,7 +135,7 @@ def start(args: Namespace):
         provider = None
         if args.provider:
             provider = LLM.load_llm(model_name=args.provider)
-        provider = provider() if provider else Cohere()
+        provider = provider() if provider else Clarifai()
         
         
         if args.api_key:

{cognitrix-0.2.4 → cognitrix-0.2.5}/cognitrix/config.py

@@ -1,9 +1,8 @@
-from pathlib import Path
 import asyncio
 import aiofiles
-import json
+from pathlib import Path
 
-VERSION = '0.2.4'
+VERSION = '0.2.5'
 SPIRAL_WORKDIR = Path('~').expanduser() / '.cognitrix'
 
 AGENTS_FILE = SPIRAL_WORKDIR / 'agents.json'

{cognitrix-0.2.4 → cognitrix-0.2.5}/cognitrix/llms/__init__.py

@@ -1,4 +1,5 @@
 from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLMResponse
 from cognitrix.llms.groq_llm import Groq
 from cognitrix.llms.local_llm import Local
 from cognitrix.llms.cohere_llm import Cohere