Jarian Cottingham
f4b84cc412
- Replace eval() with json.loads() in database.py (RCE fix) - Use json.dumps() for safe storage of list fields - Add API key authentication middleware - Remove hardcoded credentials, require env vars - Disable Flask debug mode - Restrict FTP homedir to /app/data with read-only perms - Fix threading: Lock -> RLock, add WAL mode - Fix API calls to use correct DatabaseManager methods - Fix main.py FTP method names - Fix click.click.echo typo - Implement scheduler _run_all_jobs - Add __main__.py for module execution - Pin dependency versions - Use .env vars in docker-compose, read-only DB for FTP - Implement AI text chunking with overlap windows - Add schema validation for AI responses - Skip unsupported file types instead of fallback
133 lines
4.9 KiB
Python
133 lines
4.9 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 with chunking support"""
|
|
chunk_size = 8000
|
|
overlap = 500
|
|
|
|
# Default prompt
|
|
default_prompt = """Extract key facts from the following article in structured JSON format.
|
|
Return only valid JSON without any additional text.
|
|
|
|
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"]
|
|
}"""
|
|
|
|
final_prompt = prompt if prompt else default_prompt
|
|
|
|
# Split long content into overlapping chunks
|
|
chunks = []
|
|
if len(text_content) > chunk_size:
|
|
start = 0
|
|
while start < len(text_content):
|
|
end = min(start + chunk_size, len(text_content))
|
|
chunks.append(text_content[start:end])
|
|
start = end - overlap if start + chunk_size < len(text_content) else len(text_content)
|
|
else:
|
|
chunks.append(text_content)
|
|
|
|
all_entities = set()
|
|
all_dates = set()
|
|
all_facts = []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": f"{final_prompt}\n\nArticle Content (part {i + 1}/{len(chunks)}):\n{chunk}"
|
|
}
|
|
],
|
|
"temperature": 0.3,
|
|
"max_tokens": 1000
|
|
}
|
|
|
|
try:
|
|
response = self.send_request(payload)
|
|
if 'choices' in response and len(response['choices']) > 0:
|
|
response_text = response['choices'][0]['message']['content'].strip()
|
|
if response_text.startswith('```json'):
|
|
response_text = response_text[7:-3].strip()
|
|
elif response_text.startswith('```'):
|
|
response_text = response_text[3:-3].strip()
|
|
result = json.loads(response_text)
|
|
if isinstance(result.get('key_entities'), list):
|
|
all_entities.update(result['key_entities'])
|
|
if isinstance(result.get('key_dates'), list):
|
|
all_dates.update(result['key_dates'])
|
|
if result.get('fact'):
|
|
all_facts.append(result['fact'])
|
|
except Exception as e:
|
|
print(f"Warning: Failed to process chunk {i + 1}: {str(e)}")
|
|
continue
|
|
|
|
return {
|
|
"fact": ". ".join(all_facts) if all_facts else "No facts extracted",
|
|
"key_entities": list(all_entities),
|
|
"key_dates": list(all_dates)
|
|
}
|
|
|
|
|
|
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 |