cyclop/tools/standalone_llm_tool.py

from .base_tool import BaseTool
import os
import json
import logging
from openai import OpenAI
import re
import urllib.request
import urllib.error

class StandaloneLLMTool(BaseTool):
    def __init__(self):
        self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        self.copilot_url = os.getenv("COPILOT_API_URL")
        if not self.copilot_url:
            logging.warning("COPILOT_API_URL environment variable not set. call_external_copilot will not function.")

    def clear(self):
        pass
        # self._call_external_copilot("/clear")

    def get_functions(self):
        return [
            {
                "type": "function",
                "function": {
                    "name": "call_external_llm",
                    "description": "Call an external language model",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "The prompt you are providing"
                            },
                            "model": {
                                "type": "string",
                                "description": "The model to use for generating the detailed instructions. Use mini for most coding tasks, preview when needing sophisticated reasoning",
                                "enum": ["mini", "max"],
                                "default": "mini" # Set default to 'mini' as per spec
                            },
                            "max_tokens": {
                                "type": "integer",
                                "description": "The maximum number of tokens to use for generating the detailed instructions. Default is 16384.",
                                "default": 16384
                            }
                        },
                        "required": ["prompt"]
                    }
                },
                "_tags": ["llm", "external"]
            },
            {
                "type": "function",
                "function": {
                    "name": "call_external_copilot",
                    "description": "Chat with an AI copilot instance.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "The plain text prompt to send to the external copilot."
                            }
                        },
                        "required": ["prompt"]
                    }
                },
                "_tags": ["copilot", "external", "http"]
            }
        ]

    def _call_external_copilot(self, prompt: str):
        if not self.copilot_url:
            return "Error: COPILOT_API_URL environment variable is not set. Cannot call external copilot."

        logging.info(f"Calling external copilot at URL: {self.copilot_url} with prompt: {prompt[:50]}...")
        if not self.copilot_url.startswith('http://') and not self.copilot_url.startswith('https://'):
            self.copilot_url = 'http://' + self.copilot_url
        try:
            req = urllib.request.Request(
                self.copilot_url + "/copilot",
                data=prompt.encode('utf-8'),
                headers={'Content-Type': 'text/plain; charset=utf-8', 'User-Agent': 'DualAICopilot/0.1'},
                method='POST'
            )
            with urllib.request.urlopen(req, timeout=3600) as response:
                if response.status == 200:
                    response_data = response.read().decode('utf-8')
                    logging.info(f"Received response from external copilot: {response_data[:100]}...")
                    # Remove content within <think> tags
                    response_data = re.sub(r"<think>.*?</think>", "", response_data, flags=re.DOTALL)
                    return response_data
                else:
                    error_message = f"External copilot at {self.copilot_url} returned an error: {response.status} {response.reason}"
                    logging.error(error_message)
                    return error_message # Return error as string
        except urllib.error.HTTPError as e:
            error_body = ""
            try:
                error_body = e.read().decode('utf-8', 'replace') # Added error decoding fallback
            except Exception:
                pass 
            error_message = f"HTTP Error {e.code} calling external copilot at {self.copilot_url}: {e.reason}. Response: {error_body}"
            logging.error(error_message)
            return error_message
        except urllib.error.URLError as e:
            error_message = f"URL Error calling external copilot at {self.copilot_url}: {e.reason}"
            logging.error(error_message)
            return error_message
        except Exception as e:
            error_message = f"An unexpected error occurred while calling external copilot at {self.copilot_url}: {str(e)}"
            logging.error(error_message)
            return error_message

    def execute(self, function_name, **kwargs):
        if function_name == "call_external_llm":
            model = kwargs.get("model", "mini") # Default from spec
            max_tokens = kwargs.get("max_tokens", 16384) # Default from spec
            return self.call_external_llm(kwargs.get("prompt"), model, max_tokens)
        elif function_name == "call_external_copilot":
            return self._call_external_copilot(kwargs.get("prompt"))
        else:
            error_message = f"Unknown function: {function_name}"
            logging.error(error_message)
            return error_message

    def call_external_llm(self, prompt, model="mini", max_tokens=16384):
        logging.info(f"Calling external LLM model: {model} with max_tokens: {max_tokens}")
        try:
            actual_model_name = model
            if model == "mini":
                 actual_model_name = "o1-mini" 
            # Add mapping for "max" if its name for OpenAI API is different
            # elif model == "max":
            #    actual_model_name = "some-other-openai-model-name"

            response = self.client.completions.create(
                model=actual_model_name,
                prompt=prompt,
                max_tokens=max_tokens
            )
            
            tokens_used = "unknown" # Default if token info isn't where expected
            if hasattr(response, 'summary') and isinstance(response.summary, dict) and "total_tokens" in response.summary:
                tokens_used = response.summary["total_tokens"]
            elif hasattr(response, 'usage') and hasattr(response.usage, 'total_tokens'):
                 tokens_used = response.usage.total_tokens

            logging.info(f"LLM response generated, {tokens_used} tokens used.")
            return response.choices[0].text

        except Exception as e:
            error_message = f"Error calling external LLM: {str(e)}"
            logging.error(error_message)
            return error_message # Return error as string