148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
|
|
import json
|
|
import re
|
|
import time
|
|
from typing import Optional
|
|
|
|
from .models import AgentConfig
|
|
|
|
|
|
class AgentService:
|
|
|
|
@classmethod
|
|
def execute_prompt(
|
|
cls,
|
|
config: AgentConfig,
|
|
prompt: str,
|
|
history: Optional[list] = None,
|
|
system_prompt_override: Optional[str] = None,
|
|
temperature: Optional[float] = None,
|
|
max_tokens: Optional[int] = None,
|
|
):
|
|
if not config or not config.is_active:
|
|
return False, "Agent not found or inactive", None
|
|
|
|
system = system_prompt_override or config.system_prompt or ""
|
|
temp = temperature if temperature is not None else config.temperature
|
|
tokens = max_tokens if max_tokens is not None else config.max_tokens
|
|
|
|
messages = [{"role": "system", "content": system}]
|
|
if history:
|
|
for msg in history[-20:]:
|
|
messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
try:
|
|
if config.provider == AgentConfig.Provider.OPENAI:
|
|
response = cls._call_openai(config, messages, temp, tokens)
|
|
elif config.provider == AgentConfig.Provider.DEEPSEEK:
|
|
response = cls._call_deepseek(config, messages, temp, tokens)
|
|
elif config.provider == AgentConfig.Provider.GOOGLE:
|
|
response = cls._call_google(config, messages, temp, tokens)
|
|
elif config.provider == AgentConfig.Provider.OLLAMA:
|
|
response = cls._call_ollama(config, messages, temp, tokens)
|
|
elif config.provider == AgentConfig.Provider.CUSTOM:
|
|
response = cls._call_custom(config, messages, temp, tokens)
|
|
else:
|
|
return False, f"Unsupported provider: {config.provider}", None
|
|
|
|
return True, response.get("content", ""), response.get("tokens_used")
|
|
|
|
except Exception as e:
|
|
return False, f"Error calling {config.provider}: {str(e)}", None
|
|
|
|
@classmethod
|
|
def _call_openai(cls, config, messages, temperature, max_tokens):
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=config.api_key or None)
|
|
resp = client.chat.completions.create(
|
|
model=config.model_name,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
)
|
|
return {
|
|
"content": resp.choices[0].message.content or "",
|
|
"tokens_used": resp.usage.total_tokens if resp.usage else 0,
|
|
}
|
|
|
|
@classmethod
|
|
def _call_deepseek(cls, config, messages, temperature, max_tokens):
|
|
from openai import OpenAI
|
|
client = OpenAI(
|
|
api_key=config.api_key or None,
|
|
base_url=config.base_url or "https://api.deepseek.com",
|
|
)
|
|
resp = client.chat.completions.create(
|
|
model=config.model_name or "deepseek-chat",
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
)
|
|
return {
|
|
"content": resp.choices[0].message.content or "",
|
|
"tokens_used": resp.usage.total_tokens if resp.usage else 0,
|
|
}
|
|
|
|
@classmethod
|
|
def _call_google(cls, config, messages, temperature, max_tokens):
|
|
from google import genai
|
|
client = genai.Client(api_key=config.api_key or None)
|
|
system_msg = ""
|
|
chat_messages = []
|
|
for m in messages:
|
|
if m["role"] == "system":
|
|
system_msg += m["content"] + "\n"
|
|
else:
|
|
chat_messages.append({"role": m["role"], "parts": [m["content"]]})
|
|
model = client.models.generate_content(
|
|
model=config.model_name or "gemini-2.0-flash",
|
|
contents=chat_messages,
|
|
config={
|
|
"system_instruction": system_msg.strip() if system_msg else None,
|
|
"temperature": temperature,
|
|
"max_output_tokens": max_tokens,
|
|
},
|
|
)
|
|
return {
|
|
"content": model.text or "",
|
|
"tokens_used": 0,
|
|
}
|
|
|
|
@classmethod
|
|
def _call_ollama(cls, config, messages, temperature, max_tokens):
|
|
import requests
|
|
url = (config.base_url or "http://localhost:11434") + "/api/chat"
|
|
payload = {
|
|
"model": config.model_name or "llama3",
|
|
"messages": messages,
|
|
"options": {
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens,
|
|
},
|
|
}
|
|
resp = requests.post(url, json=payload, timeout=120)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
return {
|
|
"content": data.get("message", {}).get("content", ""),
|
|
"tokens_used": 0,
|
|
}
|
|
|
|
@classmethod
|
|
def _call_custom(cls, config, messages, temperature, max_tokens):
|
|
from openai import OpenAI
|
|
client = OpenAI(
|
|
api_key=config.api_key or "fake-key",
|
|
base_url=config.base_url or "http://localhost:8000/v1",
|
|
)
|
|
resp = client.chat.completions.create(
|
|
model=config.model_name or "custom-model",
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
)
|
|
return {
|
|
"content": resp.choices[0].message.content or "",
|
|
"tokens_used": resp.usage.total_tokens if resp.usage else 0,
|
|
}
|