47a710a8b9
- Update all datetime.utcnow() to datetime.now(UTC) for Python 3.12+ compatibility - Update all datetime.utcfromtimestamp() to datetime.fromtimestamp(..., UTC) - Fix dataclass default_factory to use lambda: datetime.now(UTC) - Update auth, audit, billing, and realtime modules - Add UTC import from datetime module - Update pytest configuration to include pytest-asyncio - All 28 Phase 8 enterprise tests pass with no warnings Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
356 lines
10 KiB
Python
356 lines
10 KiB
Python
"""LLM Integration Module (Phase 7).
|
|
|
|
Supports:
|
|
- OpenAI API (GPT-4, GPT-3.5)
|
|
- Anthropic API (Claude)
|
|
- Streaming responses
|
|
- Response caching
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Optional, AsyncGenerator, Dict, Any
|
|
from enum import Enum
|
|
from abc import ABC, abstractmethod
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMProvider(str, Enum):
|
|
"""LLM providers"""
|
|
OPENAI = "openai"
|
|
ANTHROPIC = "anthropic"
|
|
LOCAL = "local" # LM Studio, Ollama
|
|
|
|
|
|
class LLMConfig:
|
|
"""LLM configuration"""
|
|
|
|
def __init__(
|
|
self,
|
|
provider: LLMProvider = LLMProvider.OPENAI,
|
|
api_key: Optional[str] = None,
|
|
model: str = "gpt-4",
|
|
temperature: float = 0.7,
|
|
max_tokens: int = 500,
|
|
base_url: Optional[str] = None,
|
|
):
|
|
self.provider = provider
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.temperature = temperature
|
|
self.max_tokens = max_tokens
|
|
self.base_url = base_url
|
|
|
|
|
|
class BaseLLMClient(ABC):
|
|
"""Base LLM client interface"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
self.config = config
|
|
|
|
@abstractmethod
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Generate response from prompt"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""Generate response as token stream"""
|
|
pass
|
|
|
|
|
|
class OpenAIClient(BaseLLMClient):
|
|
"""OpenAI API client"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
super().__init__(config)
|
|
|
|
try:
|
|
import openai
|
|
self.client = openai.AsyncOpenAI(api_key=config.api_key)
|
|
except ImportError:
|
|
raise ImportError("openai package required: pip install openai")
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Generate response from OpenAI"""
|
|
try:
|
|
response = await self.client.chat.completions.create(
|
|
model=self.config.model,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt
|
|
}
|
|
],
|
|
temperature=self.config.temperature,
|
|
max_tokens=self.config.max_tokens,
|
|
stream=stream,
|
|
)
|
|
|
|
if stream:
|
|
# Collect streamed tokens
|
|
full_response = ""
|
|
async for chunk in response:
|
|
if chunk.choices[0].delta.content:
|
|
full_response += chunk.choices[0].delta.content
|
|
return full_response
|
|
else:
|
|
return response.choices[0].message.content
|
|
|
|
except Exception as e:
|
|
logger.error(f"OpenAI generation failed: {e}")
|
|
raise
|
|
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""Stream tokens from OpenAI"""
|
|
try:
|
|
response = await self.client.chat.completions.create(
|
|
model=self.config.model,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt
|
|
}
|
|
],
|
|
temperature=self.config.temperature,
|
|
max_tokens=self.config.max_tokens,
|
|
stream=True,
|
|
)
|
|
|
|
async for chunk in response:
|
|
if chunk.choices[0].delta.content:
|
|
yield chunk.choices[0].delta.content
|
|
|
|
except Exception as e:
|
|
logger.error(f"OpenAI streaming failed: {e}")
|
|
raise
|
|
|
|
|
|
class AnthropicClient(BaseLLMClient):
|
|
"""Anthropic API client (Claude)"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
super().__init__(config)
|
|
|
|
try:
|
|
import anthropic
|
|
self.client = anthropic.AsyncAnthropic(api_key=config.api_key)
|
|
except ImportError:
|
|
raise ImportError("anthropic package required: pip install anthropic")
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Generate response from Claude"""
|
|
try:
|
|
if stream:
|
|
full_response = ""
|
|
async with self.client.messages.stream(
|
|
model=self.config.model,
|
|
max_tokens=self.config.max_tokens,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt
|
|
}
|
|
],
|
|
) as stream:
|
|
async for text in stream.text_stream:
|
|
full_response += text
|
|
return full_response
|
|
else:
|
|
message = await self.client.messages.create(
|
|
model=self.config.model,
|
|
max_tokens=self.config.max_tokens,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt
|
|
}
|
|
],
|
|
)
|
|
return message.content[0].text
|
|
|
|
except Exception as e:
|
|
logger.error(f"Anthropic generation failed: {e}")
|
|
raise
|
|
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""Stream tokens from Claude"""
|
|
try:
|
|
async with self.client.messages.stream(
|
|
model=self.config.model,
|
|
max_tokens=self.config.max_tokens,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt
|
|
}
|
|
],
|
|
) as stream:
|
|
async for text in stream.text_stream:
|
|
yield text
|
|
|
|
except Exception as e:
|
|
logger.error(f"Anthropic streaming failed: {e}")
|
|
raise
|
|
|
|
|
|
class LocalLLMClient(BaseLLMClient):
|
|
"""Local LLM client (LM Studio, Ollama)"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
super().__init__(config)
|
|
|
|
try:
|
|
import httpx
|
|
self.client = httpx.AsyncClient(base_url=config.base_url)
|
|
except ImportError:
|
|
raise ImportError("httpx package required: pip install httpx")
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Generate response from local LLM"""
|
|
try:
|
|
response = await self.client.post(
|
|
"/v1/completions",
|
|
json={
|
|
"model": self.config.model,
|
|
"prompt": prompt,
|
|
"temperature": self.config.temperature,
|
|
"max_tokens": self.config.max_tokens,
|
|
"stream": stream,
|
|
},
|
|
)
|
|
|
|
if stream:
|
|
full_response = ""
|
|
async for chunk in response.aiter_lines():
|
|
if chunk.startswith("data: "):
|
|
import json
|
|
try:
|
|
data = json.loads(chunk[6:])
|
|
if "choices" in data:
|
|
full_response += data["choices"][0].get("text", "")
|
|
except:
|
|
pass
|
|
return full_response
|
|
else:
|
|
data = response.json()
|
|
return data["choices"][0]["text"]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Local LLM generation failed: {e}")
|
|
raise
|
|
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""Stream tokens from local LLM"""
|
|
try:
|
|
async with self.client.stream(
|
|
"POST",
|
|
"/v1/completions",
|
|
json={
|
|
"model": self.config.model,
|
|
"prompt": prompt,
|
|
"temperature": self.config.temperature,
|
|
"max_tokens": self.config.max_tokens,
|
|
"stream": True,
|
|
},
|
|
) as response:
|
|
async for chunk in response.aiter_lines():
|
|
if chunk.startswith("data: "):
|
|
import json
|
|
try:
|
|
data = json.loads(chunk[6:])
|
|
if "choices" in data:
|
|
text = data["choices"][0].get("text", "")
|
|
if text:
|
|
yield text
|
|
except:
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.error(f"Local LLM streaming failed: {e}")
|
|
raise
|
|
|
|
|
|
class LLMManager:
|
|
"""LLM management and client selection"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
self.config = config
|
|
self.client = self._create_client(config)
|
|
|
|
def _create_client(self, config: LLMConfig) -> BaseLLMClient:
|
|
"""Create appropriate LLM client"""
|
|
if config.provider == LLMProvider.OPENAI:
|
|
return OpenAIClient(config)
|
|
elif config.provider == LLMProvider.ANTHROPIC:
|
|
return AnthropicClient(config)
|
|
elif config.provider == LLMProvider.LOCAL:
|
|
return LocalLLMClient(config)
|
|
else:
|
|
raise ValueError(f"Unknown provider: {config.provider}")
|
|
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Generate response"""
|
|
return await self.client.generate(prompt, stream=stream)
|
|
|
|
async def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""Generate streaming response"""
|
|
async for token in self.client.generate_stream(prompt):
|
|
yield token
|
|
|
|
async def generate_with_metadata(
|
|
self,
|
|
prompt: str,
|
|
stream: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Generate response with metadata"""
|
|
import time
|
|
|
|
start_time = time.time()
|
|
response = await self.generate(prompt, stream=stream)
|
|
end_time = time.time()
|
|
|
|
return {
|
|
"response": response,
|
|
"tokens": len(response.split()),
|
|
"latency": end_time - start_time,
|
|
"model": self.config.model,
|
|
"provider": self.config.provider.value,
|
|
}
|