156 lines
7.1 KiB
Python
156 lines
7.1 KiB
Python
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
|