303 lines
11 KiB
Python
303 lines
11 KiB
Python
import os
|
|
import json
|
|
import chromadb
|
|
import uuid
|
|
import time
|
|
import datetime
|
|
import requests
|
|
import logging
|
|
from pathlib import Path
|
|
import openai
|
|
from prometheus_client import start_http_server, Counter, Histogram
|
|
|
|
# Setup logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Start Prometheus metrics server
|
|
try:
|
|
start_http_server(8001)
|
|
logger.info("Prometheus metrics server started on port 8001")
|
|
except Exception as e:
|
|
logger.error(f"Failed to start Prometheus server: {e}")
|
|
|
|
# Prometheus metrics for the embedding pipeline
|
|
articles_processed_total = Counter('embedding_pipeline_articles_processed_total', 'Total number of articles processed')
|
|
articles_failed_total = Counter('embedding_pipeline_articles_failed_total', 'Total number of articles failed to process')
|
|
processing_time_seconds = Histogram('embedding_pipeline_processing_time_seconds', 'Time spent processing articles')
|
|
|
|
# ChromaDB client setup
|
|
CHROMADB_HOST = os.getenv("CHROMADB_HOST", "example.com")
|
|
CHROMADB_PORT = int(os.getenv("CHROMADB_PORT", "8000"))
|
|
|
|
try:
|
|
client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
|
|
logger.info("Connected to ChromaDB successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to ChromaDB: {e}")
|
|
raise
|
|
|
|
# AI Server configuration
|
|
AI_SERVER_HOST = os.getenv("AI_SERVER_HOST", "example.com")
|
|
AI_SERVER_PORT = int(os.getenv("AI_SERVER_PORT", "4000"))
|
|
AI_SERVER_URL = f"http://{AI_SERVER_HOST}:{AI_SERVER_PORT}/v1/embeddings"
|
|
|
|
# Cache file for tracking processed articles
|
|
CACHE_FILE = os.getenv("CACHE_FILE", "processed_articles_cache.json")
|
|
|
|
# OpenAI client for fact extraction (if using local LLM)
|
|
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "placeholder-key")
|
|
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", f"http://{AI_SERVER_HOST}:{AI_SERVER_PORT}")
|
|
|
|
def get_embedding(text):
|
|
"""
|
|
Get embedding using the OpenAI-compatible server
|
|
"""
|
|
try:
|
|
response = requests.post(
|
|
AI_SERVER_URL,
|
|
json={
|
|
"input": text,
|
|
"model": "qwen3:8b"
|
|
},
|
|
timeout=30
|
|
)
|
|
response.raise_for_status()
|
|
embedding = response.json()['data'][0]['embedding']
|
|
return embedding
|
|
except Exception as e:
|
|
logger.error(f"Error getting embedding: {e}")
|
|
return None
|
|
|
|
def extract_facts_from_article(article_content, title):
|
|
"""
|
|
Extract structured facts from article content using LLM with proper prompting
|
|
"""
|
|
try:
|
|
# Create a proper prompt for fact extraction
|
|
prompt = f"""
|
|
Extract key facts from the following article in structured JSON format.
|
|
Return only valid JSON without any additional text.
|
|
|
|
Article Title: {title}
|
|
Article Content: {article_content[:1000]}...
|
|
|
|
Extract the following information:
|
|
1. Main topic/subject
|
|
2. Key entities (companies, people, locations, organizations)
|
|
3. Financial impact or implications
|
|
4. Key dates or time periods mentioned
|
|
5. Summary of main points
|
|
|
|
Format the response as a JSON object with these fields:
|
|
{{
|
|
"title": "article title",
|
|
"summary": "brief summary",
|
|
"main_topic": "main topic",
|
|
"key_entities": ["entity1", "entity2"],
|
|
"financial_impact": "positive/negative/neutral",
|
|
"key_dates": ["date1", "date2"],
|
|
"main_points": ["point1", "point2", "point3"]
|
|
}}
|
|
"""
|
|
|
|
# For now, using the existing embedding approach - in a real implementation
|
|
# this would call the AI server with a proper prompt
|
|
# response = requests.post(AI_SERVER_URL, json={
|
|
# "model": "gpt-4",
|
|
# "messages": [
|
|
# {"role": "system", "content": "You are a helpful assistant that extracts structured facts from articles."},
|
|
# {"role": "user", "content": prompt}
|
|
# ]
|
|
# })
|
|
|
|
# Simplified version for now - in production this would be a proper LLM call
|
|
facts = {
|
|
"title": title,
|
|
"summary": article_content[:200] + "..." if len(article_content) > 200 else article_content,
|
|
"main_topic": "Business/Financial News",
|
|
"key_entities": ["Sample Corp", "John Doe"],
|
|
"financial_impact": "neutral",
|
|
"key_dates": ["2026"],
|
|
"main_points": [
|
|
"This is a sample key point extracted from the article",
|
|
"Another important fact from the content",
|
|
"Third key fact from the article"
|
|
]
|
|
}
|
|
|
|
return facts
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error extracting facts: {e}")
|
|
# Return a basic structure if extraction fails
|
|
return {
|
|
"title": title,
|
|
"summary": article_content[:200] + "..." if len(article_content) > 200 else article_content,
|
|
"main_topic": "Unknown",
|
|
"key_entities": [],
|
|
"financial_impact": "neutral",
|
|
"key_dates": [],
|
|
"main_points": []
|
|
}
|
|
|
|
def process_article_file(file_path):
|
|
"""
|
|
Process a single article file and extract facts
|
|
"""
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
article_data = json.load(f)
|
|
|
|
# Extract facts from the article
|
|
facts = extract_facts_from_article(
|
|
article_data.get('original_content', ''),
|
|
article_data.get('title', '')
|
|
)
|
|
|
|
# Add metadata
|
|
facts['source'] = article_data.get('source', 'Unknown')
|
|
facts['published'] = article_data.get('published', 'Unknown')
|
|
facts['filename'] = os.path.basename(file_path)
|
|
facts['processed_at'] = datetime.datetime.now().isoformat()
|
|
|
|
return facts
|
|
except Exception as e:
|
|
logger.error(f"Error processing article {file_path}: {e}")
|
|
return None
|
|
|
|
def create_collections():
|
|
"""
|
|
Create necessary ChromaDB collections for different types of data
|
|
"""
|
|
# Collection for extracted facts (now with entity support)
|
|
facts_collection = client.get_or_create_collection("facts")
|
|
|
|
# Collection for full articles
|
|
articles_collection = client.get_or_create_collection("articles")
|
|
|
|
# Remove company collection - now using entity tracking in facts collection
|
|
return facts_collection, articles_collection
|
|
|
|
def embed_and_store_facts(facts, facts_collection, articles_collection):
|
|
"""
|
|
Embed and store facts in appropriate collections
|
|
"""
|
|
try:
|
|
# Store the complete article in articles collection
|
|
article_embedding = get_embedding(facts['title'] + " " + facts['summary'])
|
|
if article_embedding:
|
|
articles_collection.upsert(
|
|
ids=[str(uuid.uuid4())],
|
|
documents=[facts['title'] + " " + facts['summary']],
|
|
embeddings=[article_embedding],
|
|
metadatas=[{
|
|
"source": facts['source'],
|
|
"published": facts['published'],
|
|
"filename": facts['filename'],
|
|
"type": "article",
|
|
"processed_at": facts['processed_at'],
|
|
"main_topic": facts.get('main_topic', 'Unknown')
|
|
}]
|
|
)
|
|
|
|
# Store extracted facts in facts collection
|
|
facts_text = json.dumps(facts, indent=2)
|
|
facts_embedding = get_embedding(facts_text)
|
|
if facts_embedding:
|
|
facts_collection.upsert(
|
|
ids=[str(uuid.uuid4())],
|
|
documents=[facts_text],
|
|
embeddings=[facts_embedding],
|
|
metadatas=[{
|
|
"source": facts['source'],
|
|
"published": facts['published'],
|
|
"filename": facts['filename'],
|
|
"type": "fact",
|
|
"processed_at": facts['processed_at'],
|
|
"title": facts['title'],
|
|
"main_topic": facts.get('main_topic', 'Unknown'),
|
|
"key_entities": facts.get('key_entities', []),
|
|
"financial_impact": facts.get('financial_impact', 'neutral')
|
|
}]
|
|
)
|
|
|
|
logger.info(f"Successfully processed and stored facts for {facts['filename']}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error embedding and storing facts: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""
|
|
Main embedding pipeline function
|
|
"""
|
|
logger.info("Starting advanced embedding pipeline")
|
|
|
|
# Create collections
|
|
facts_collection, articles_collection = create_collections()
|
|
|
|
# Load cache of previously processed articles
|
|
processed_cache = {}
|
|
if os.path.exists(CACHE_FILE):
|
|
with open(CACHE_FILE, 'r', encoding='utf-8') as f:
|
|
processed_cache = json.load(f)
|
|
|
|
# Process articles from scraper directory
|
|
scraper_articles_dir = "/scraper/articles"
|
|
|
|
# Track processing time
|
|
start_time = datetime.datetime.now()
|
|
|
|
# Walk through all subdirectories in scraper articles
|
|
for root, dirs, files in os.walk(scraper_articles_dir):
|
|
for file in files:
|
|
if file.endswith('.json'):
|
|
file_path = os.path.join(root, file)
|
|
|
|
# Check if already processed
|
|
if file_path in processed_cache:
|
|
logger.info(f"Article {file} already processed, skipping.")
|
|
continue
|
|
|
|
# Process the article
|
|
facts = process_article_file(file_path)
|
|
if facts:
|
|
# Embed and store in appropriate collections
|
|
success = embed_and_store_facts(
|
|
facts,
|
|
facts_collection,
|
|
articles_collection
|
|
)
|
|
|
|
if success:
|
|
processed_cache[file_path] = {
|
|
"processed_date": datetime.datetime.now().isoformat(),
|
|
"status": "completed"
|
|
}
|
|
articles_processed_total.inc()
|
|
logger.info(f"Successfully processed {file}")
|
|
else:
|
|
articles_failed_total.inc()
|
|
logger.error(f"Failed to process {file}")
|
|
|
|
# Save updated cache
|
|
try:
|
|
with open(CACHE_FILE, 'w', encoding='utf-8') as f:
|
|
json.dump(processed_cache, f, indent=2)
|
|
logger.info(f"Updated cache with newly processed articles. Total cached: {len(processed_cache)}")
|
|
except Exception as e:
|
|
logger.error(f"Error saving cache file: {e}")
|
|
|
|
# Calculate and log processing time
|
|
end_time = datetime.datetime.now()
|
|
total_time = (end_time - start_time).total_seconds()
|
|
processing_time_seconds.observe(total_time)
|
|
logger.info(f"Embedding pipeline completed in {total_time:.2f} seconds")
|
|
|
|
logger.info("Embedding pipeline completed")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|