127 lines
4.8 KiB
Python
127 lines
4.8 KiB
Python
"""
|
|
AI processor module for FactsDB service
|
|
Handles communication with OpenAI compatible endpoints
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
from typing import Dict, Any, Optional
|
|
from .config import AIEndpointConfig
|
|
|
|
class AIEndpointClient:
|
|
"""Client for communicating with OpenAI compatible endpoint"""
|
|
|
|
def __init__(self, config: AIEndpointConfig):
|
|
self.config = config
|
|
self.base_url = config.url.rstrip('/')
|
|
self.auth_token = config.auth_token
|
|
|
|
def _get_headers(self) -> Dict[str, str]:
|
|
"""Get headers with authentication"""
|
|
return {
|
|
'Authorization': f'Bearer {self.auth_token}',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
def send_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Send request to AI endpoint"""
|
|
url = f"{self.base_url}/v1/chat/completions"
|
|
|
|
try:
|
|
response = requests.post(
|
|
url,
|
|
headers=self._get_headers(),
|
|
json=payload,
|
|
timeout=300 # 5 minute timeout
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except requests.exceptions.RequestException as e:
|
|
raise Exception(f"AI endpoint request failed: {str(e)}")
|
|
|
|
def extract_facts(self, text_content: str, prompt: str, model: str = "gpt-oss") -> Dict[str, Any]:
|
|
"""Extract facts from text using AI"""
|
|
# Default prompt from requirements - optimized for facts extraction
|
|
default_prompt = """Extract key facts from the following article in structured JSON format.
|
|
Return only valid JSON without any additional text.
|
|
|
|
Article Content: {article_content[:3000]}...
|
|
|
|
Extract the following information:
|
|
1. Key entities (companies, people, locations, organizations)
|
|
2. Key dates or time periods mentioned
|
|
3. Main facts from the article
|
|
|
|
Format the response as a JSON object with these fields:
|
|
{
|
|
"fact": "main fact extracted from the article",
|
|
"key_entities": ["entity1", "entity2"],
|
|
"key_dates": ["date1", "date2"]
|
|
}"""
|
|
|
|
# Use provided prompt or default
|
|
final_prompt = prompt if prompt else default_prompt
|
|
|
|
# Create the payload
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": f"{final_prompt}\n\nArticle Content: {text_content[:3000]}"
|
|
}
|
|
],
|
|
"temperature": 0.3,
|
|
"max_tokens": 1000
|
|
}
|
|
|
|
try:
|
|
response = self.send_request(payload)
|
|
|
|
# Extract the response text
|
|
if 'choices' in response and len(response['choices']) > 0:
|
|
response_text = response['choices'][0]['message']['content']
|
|
|
|
# Try to parse JSON
|
|
try:
|
|
# Clean up the response to ensure valid JSON
|
|
response_text = response_text.strip()
|
|
if response_text.startswith('```json'):
|
|
response_text = response_text[7:-3].strip()
|
|
elif response_text.startswith('```'):
|
|
response_text = response_text[3:-3].strip()
|
|
|
|
return json.loads(response_text)
|
|
except json.JSONDecodeError:
|
|
# If JSON parsing fails, return the raw response as a structured format
|
|
return {
|
|
"raw_response": response_text,
|
|
"title": "Unknown",
|
|
"summary": response_text[:200] + "..." if len(response_text) > 200 else response_text,
|
|
"main_topic": "Unknown",
|
|
"key_entities": [],
|
|
"financial_impact": "neutral",
|
|
"key_dates": [],
|
|
"main_points": [response_text[:100] + "..."] if len(response_text) > 100 else [response_text]
|
|
}
|
|
else:
|
|
raise Exception("No response from AI model")
|
|
|
|
except Exception as e:
|
|
raise Exception(f"Fact extraction failed: {str(e)}")
|
|
|
|
class AIProcessor:
|
|
"""Main AI processor class for FactsDB"""
|
|
|
|
def __init__(self, config: AIEndpointConfig):
|
|
self.client = AIEndpointClient(config)
|
|
|
|
def extract_facts_from_text(self, text_content: str, prompt: str = "", model: str = "gpt-oss") -> Dict[str, Any]:
|
|
"""Extract facts from text content using AI"""
|
|
return self.client.extract_facts(text_content, prompt, model)
|
|
|
|
def validate_model_support(self, model: str) -> bool:
|
|
"""Validate if model is supported"""
|
|
# In a real implementation, this would check against available models
|
|
supported_models = ["gpt-oss", "qwen3", "qwen3-coder"]
|
|
return model in supported_models |