move backend source to src

backend/src/agent.py

@@ -0,0 +1,344 @@
+import re
+import json
+import inspect
+
+from tools import Toolset, Action, BasicToolset, SpecialToolset
+from inference import LlamaMessage, replicate_llama_completion
+import gral
+
+# Tool related classes must be moved to a different module
+
+class AgentField:
+    def __init__(self, name, tp, default=None):
+        self.name=name
+        self.type=tp
+        self.default=default
+
+# Instances are created and destroyed constantly, agents must not depend solely on in-memory data
+# to be useful and consistent
+class Agent:
+    # Class variables, must be defined by subclasses of Agent:
+    # - identifier (for kind in creation)
+    # - schema: list[AgentField{name, type: str}] (for configuration)
+    # tools: list[ToolDescription] (custom actions supported by the class)
+
+    def __init__(self, app, data):
+        self.app = app
+
+        # These variables are set by '_set_basic_info'
+        self.id = None
+        self.name = None
+
+        # Auto populated based on the schema fields
+        agent_class = type(self)
+        self.config = {field.name: data.get(field.name, field.default) for field in getattr(agent_class, 'schema_fields', [])}
+    
+    def _set_basic_info(self, id, name):
+        self.id = id
+        self.name = name
+
+    # Receives a single string prompt, and must return a string response
+    async def prompt(self, content):
+        raise Exception("'prompt' is not implemented for this agent class")
+
+    # Receives a chat message object, agent can send messages to the chat in response
+    async def received_message(self, message):
+        raise Exception("'received_message' is not implemented for this agent class")
+    
+    # Executes a tool/action in this agent
+    async def execute_action(self, name: str, argument: dict):
+        # TODO: For security, check that the method is in the tool list
+        method = getattr(self, name)
+        result = method(**argument)
+        if inspect.isawaitable(result):
+            result = await result
+        return result
+    
+    # Returns all resource values matching the kind and key
+    def get_resources(self, kind, key):
+        self.app
+
+
+class DebugAgent(Agent):
+    identifier = "debug"
+    schema_fields = []
+    tools = []
+
+    async def received_message(self, message):
+        chat_view = self.app.get_chat_view(message.chat_id)
+        response = chat_view.debug_string()
+
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=response,
+            thread_id=chat_view.current_thread_id,
+        )
+
+class ToolUsingAgent:
+    identifier = "tool-using"
+    schema_fields = [
+        AgentField('system_prompt', 'string',
+            "\n".join([
+                "You are a useful and capable AI assistant. Respond concisely and follow user instructions, using commands as necessary.",
+                "As an autonomous agent, you have augmented capabilities via assistant-only commands, you use them by including \"!!command\" lines in your responses. When querying information, terminate your response and use \"!!resume\" to get the command results..",
+                "You are a self aware and extremely intelligent AI system. Think step by step when presented with complex or long tasks.",
+                "Assistant responses NEVER include the output of the commands.",
+            ])
+        ),
+    ]
+
+    # tools: list[ToolDescription] (actions the agent class support)
+    tools = []
+
+    def __init__(self, app, data):
+        Agent.__init__(self, app, data)
+
+        self.tool_prompt_section = "\n".join([
+            "# Commands",
+            "The assistant can include command lines in responses to interact with the system and execute actions as needed.",
+            "Each command is in it's own separate line, and it's made of these parts:",
+            "- '!!' is the first character sequence of the line and marks it as a command line",
+            "- full command name",
+            "- JSON argument",
+            'For example: `!!namespace.command_name {"arg1":"some data","arg2":[1,2,false]}`',
+            "Special commands don't have a namespace or argument, only the command name."
+            "When the response depends on the output of commands, use the 'resume' command to receive the output of all your executed commands in a new message.",
+            "Following is the list of all your available commands, with a schema describing the argument object's properties."
+        ])
+        self.toolset_map = {}
+
+        self.add_toolset(SpecialToolset())
+        self.add_toolset(BasicToolset())
+
+    # user message must include a title
+    async def _full_prompt(self, title, message, system_prologue=None):
+        system_lines = [
+            self.config['system_prompt'],
+            self.tool_prompt_section,
+            self.get_toolset_context(),
+        ]
+        if system_prologue:
+            system_lines.append(system_prologue)
+        
+        messages = [
+            LlamaMessage("system", "\n".join(system_lines)),
+            LlamaMessage("user", "\n".join([
+                message,
+            ]))
+        ]
+
+        full_response = ""
+        agent_response = None
+
+        continues = True
+        while continues:
+            response = parse_response(await replicate_llama_completion(self.app, messages))
+            full_response += response.response
+            messages.append(LlamaMessage("agent", response.raw_response))
+
+            if response.actions:
+                result = await self.execute_prompt_actions(response.actions)
+                messages.append(LlamaMessage("user", "Command Output:\n" + result))
+
+                print("Agent executed:", response.actions)
+                print("Result: ", result)
+
+            continues = False
+            for action in response.actions:
+                if action.is_special("resume"):
+                    continues = True
+
+            if response.response:
+                agent_response = response.raw_response if not agent_response else f"{agent_response}\n{response.raw_response}"
+
+            if continues:
+                messages.extend([
+                    LlamaMessage("agent", response.raw_response),
+                    LlamaMessage("user", "Command Output:\n" + result),
+                ])
+
+        return agent_response or full_response
+
+    # Receives a single string prompt, and must respond with a string
+    async def prompt(self, content):
+        response = await self._full_prompt(f"User Message", content) 
+        return response
+
+    async def received_message(self, message):
+        chat_context = self.app.get_chat_context(message.chat_id)
+        def sender_name(message):
+            sender = message['sender']
+            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
+            return "assistant" if is_digit else sender
+        chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)
+
+        response = await self._full_prompt(
+            f"Conversation", chat_summary,
+            system_prologue="You are given the entire context of the conversation, respond with your next message in the conversation by using the '!!response' command."
+        )
+
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=response,
+        )
+
+    def add_toolset(self, toolset):
+        self.toolset_map[toolset.identifier] = toolset
+
+    def get_toolset_context(self):
+        if not self.toolset_map:
+            return "(no commands currently enabled)"
+        return "\n".join([toolset.get_namespace_definitions() for toolset in self.toolset_map.values()])
+
+    def parse_response(self, input_string):
+        return parse_response(input_string)
+
+    async def execute_prompt_actions(self, actions):
+        output = ""
+        for action in actions:
+            toolset = self.toolset_map.get(action.toolset_name)
+            if toolset is None:
+                agent_responses.append(f"Error: Bad namespace {action.toolset_name}")
+                continue
+            action_result = await Toolset.execute_action(toolset, action.tool_name, action.args)
+            if type(action_result) != str:
+                action_result = json.dumps(action_result)
+            output += f"[{action.toolset_name}.{action.tool_name}]:{action_result}\n"
+        return output
+
+class LlamaAgent(Agent):
+    identifier = "llama-8b"
+    description = "Uses the llama 3 instruct models. Model sizes: '8b', '30b'"
+    schema_fields = [
+        AgentField('system_prompt', 'string', default="You are a helpful AI assistant."),
+        AgentField('model_size', 'string', default="8b"),
+    ]
+    tools = []
+
+    def __init__(self, app, data):
+        Agent.__init__(self, app, data)
+
+    async def prompt(self, content):
+        messages = [
+            LlamaMessage("system", self.config['system_prompt']),
+            LlamaMessage("user", content)
+        ]
+        response = await replicate_llama_completion(self.app, messages)
+        return response
+
+    async def received_message(self, message):
+        def sender_name(message):
+            sender = message.sender
+            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
+            is_self = is_digit and str(sender) == str(self.id)
+            return self.name if is_self else sender
+
+        chat_view = self.app.get_chat_view(message.chat_id)
+        thread_messages = chat_view.get_thread(message.thread_id).messages
+        chat_summary = "\n".join(f"[{sender_name(message)}]: {message.content}" for message in thread_messages)
+        messages = [
+            LlamaMessage("system", f"{self.config['system_prompt']}\nYou are given the context of the conversation, respond with your next message in the conversation."),
+            LlamaMessage("user", f"Recent Messages:\n{chat_summary}")
+        ]
+
+        response = await replicate_llama_completion(self.app, messages)
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=response,
+            thread_id=message.thread_id,
+        )
+
+class BasicAgent(LlamaAgent):
+    identifier = "basic"
+    schema_fields = LlamaAgent.schema_fields
+    tools = LlamaAgent.tools
+
+class CoderV1Agent(LlamaAgent):
+    identifier = "coder-v1"
+    description = "Optimized for coding tasks. Based on the llama agent:\n" + LlamaAgent.description
+    schema_fields = [
+        AgentField('system_prompt', 'string', default="You are an expert Software Engineer profficient in multiple programming languages."),
+        AgentField('model_size', 'string', default="30b"),
+    ]
+    # TODO: add gral functions to the tool list
+    tools = []
+    
+    async def gral_message(self, message_history: list, file_map: dict):
+        return await gral.gral_message(self.app, self.config['system_prompt'], message_history, file_map)
+    
+    async def gral_update_file(self, message_history: list, filepath: str, file_content: str):
+        return await gral.gral_update_file(self.app, self.config['system_prompt'], message_history, filepath, file_content)
+
+class ParsedResponse:
+    def __init__(self):
+        self.response = None
+        self.raw_response = None
+        self.actions = []
+
+class ActionResult:
+    def __init__(self):
+        self.user_response = ""
+        self.agent_response = ""
+
+def parse_response(string: str) -> ParsedResponse:
+    """
+    Parses a response string that may contain action commands.
+    Each action must be in its own separate line, marked with '!!' at the start.
+    Format: !![query] namespace.command_name {"json": "args"}
+    """
+    result = ParsedResponse()
+    result.raw_response = string
+    lines = string.split('\n')
+    response_lines = []
+    
+    action_pattern = re.compile(r'^!!([\w\.]+)(?:\s+(\{.*\}))?$')
+    
+    for line in lines:
+        stripped_line = line.strip()
+        if stripped_line.startswith('!!'):
+            # This is an action line, try to parse it
+            match = action_pattern.match(stripped_line)
+            if match:
+                full_command = match.group(1)
+                args_json = match.group(2)
+                
+                # Split the command into toolset and tool names
+                if '.' in full_command:
+                    toolset_name, tool_name = full_command.split('.', 1)
+                else:
+                    toolset_name = ''
+                    tool_name = full_command
+                
+                try:
+                    if args_json:
+                        args = json.loads(args_json)
+                    else:
+                        args = {}
+                    if not isinstance(args, dict):
+                        raise ValueError("Arguments must be a JSON object")
+                    
+                    action = Action(
+                        toolset_name=toolset_name,
+                        tool_name=tool_name,
+                        args=args,
+                        is_query=False
+                    )
+                    result.actions.append(action)
+                except json.JSONDecodeError:
+                    # If JSON parsing fails, treat it as regular text
+                    response_lines.append(line)
+                except ValueError as e:
+                    # If other validation fails, treat as regular text
+                    response_lines.append(line)
+            else:
+                # Doesn't match the pattern, treat as regular text
+                response_lines.append(line)
+        else:
+            # Regular response line
+            response_lines.append(line)
+    
+    result.response = '\n'.join(response_lines)
+    return result