125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
from typing import List, Dict, Any, Optional, Tuple
|
|
import os
|
|
import requests
|
|
import time
|
|
import json
|
|
import logging
|
|
import traceback
|
|
from .constants import MAX_HISTORY_LENGTH, load_config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class LLMClient:
|
|
"""
|
|
Client for interacting with LLM models via KeyPool service.
|
|
Handles conversation history and API requests.
|
|
"""
|
|
|
|
def __init__(self, model_name: str, keypool_url: str=None, api_key: str=None):
|
|
"""
|
|
Initialize LLM client.
|
|
|
|
Args:
|
|
model_name: Name of the model to use (e.g., "gemini-2.5-flash-lite")
|
|
Must match a model_name in KeyPool's keys.yaml
|
|
"""
|
|
self.history: List[Dict[str, Any]] = []
|
|
|
|
# Default timeout for requests
|
|
self.timeout: int = 120
|
|
|
|
# KeyPool Service Configuration
|
|
config = load_config() if (keypool_url is None or api_key is None) else {}
|
|
if model_name is None:
|
|
model_name = (
|
|
config.get("CURRENT_MODEL_NAME")
|
|
or config.get("llm_api", {}).get("current_model_name", "")
|
|
or ""
|
|
)
|
|
self.model_name = str(model_name).strip()
|
|
if keypool_url is None:
|
|
keypool_url = (
|
|
config.get("KEY_POOL_URL")
|
|
or config.get("llm_api", {}).get("key_pool_url", "")
|
|
or ""
|
|
)
|
|
self.keypool_url = str(keypool_url).strip().rstrip("/")
|
|
|
|
if api_key is None:
|
|
api_key = (
|
|
config.get("KEY_POOL_API_KEY")
|
|
or config.get("llm_api", {}).get("key_pool_api_key", "")
|
|
or os.environ.get("KEY_POOL_API_KEY", "")
|
|
or ""
|
|
)
|
|
self.api_key = str(api_key).strip()
|
|
|
|
logger.info(f"LLMClient initialized for model: {self.model_name}")
|
|
|
|
def query(self, messages: List[Dict[str, Any]], extra_params: Optional[Dict[str, Any]] = None) -> Tuple[str, Any, Dict[str, Any]]:
|
|
"""
|
|
Send a query to the LLM via KeyPool.
|
|
|
|
Args:
|
|
messages: List of message dictionaries
|
|
extra_params: Optional dictionary of extra parameters to override defaults
|
|
|
|
Returns:
|
|
Tuple[str, Any, Dict[str, Any]]: (response_content, reasoning_content, token_usage)
|
|
"""
|
|
try:
|
|
# Prepare payload for KeyPool
|
|
payload = {
|
|
"model": self.model_name,
|
|
"messages": messages,
|
|
"stream": False
|
|
}
|
|
|
|
# Add extra parameters if provided
|
|
if extra_params:
|
|
payload.update(extra_params)
|
|
|
|
# Build headers (OpenAI-compatible /v1/chat/completions endpoint)
|
|
headers = {"Content-Type": "application/json"}
|
|
if self.api_key:
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
|
|
# Send request to KeyPool (OpenAI-compatible endpoint)
|
|
response = requests.post(
|
|
f"{self.keypool_url}/v1/chat/completions",
|
|
json=payload,
|
|
headers=headers,
|
|
timeout=self.timeout
|
|
)
|
|
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
content = ""
|
|
reasoning = None
|
|
usage = {}
|
|
|
|
if "choices" in result and len(result["choices"]) > 0:
|
|
message = result["choices"][0]["message"]
|
|
content = message.get("content", "")
|
|
reasoning = message.get("reasoning_content")
|
|
|
|
if "usage" in result:
|
|
usage = result["usage"]
|
|
return content, reasoning, usage
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error querying KeyPool: {e}")
|
|
logger.error(traceback.format_exc())
|
|
return "", None, {}
|
|
|
|
def add_history(self, role: str, content: Any):
|
|
"""Add a message to history"""
|
|
self.history.append({"role": role, "content": content})
|
|
if len(self.history) > MAX_HISTORY_LENGTH:
|
|
self.history.pop(0)
|
|
|
|
def clear_history(self):
|
|
"""Clear conversation history"""
|
|
self.history = []
|