feat: implement enhanced data pipeline with Prometheus telemetry and Grafana integration
This commit is contained in:
parent
5e52142909
commit
c8a712d776
@ -264,6 +264,53 @@
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/facts": {
|
||||
"post": {
|
||||
"summary": "Get the best set of facts about a question",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": { "type": "string" }
|
||||
},
|
||||
"required": ["question"]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Facts and articles related to the question",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": { "type": "string" },
|
||||
"articles": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"document": { "type": "string" },
|
||||
"score": { "type": "number" },
|
||||
"metadata": { "type": "object" }
|
||||
}
|
||||
}
|
||||
},
|
||||
"company_facts": { "type": "object" },
|
||||
"timestamp": { "type": "string", "format": "date-time" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
325
embedding/advanced_embedder.py
Normal file
325
embedding/advanced_embedder.py
Normal file
@ -0,0 +1,325 @@
|
||||
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": "text-embedding-3-small"
|
||||
},
|
||||
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
|
||||
facts_collection = client.get_or_create_collection("facts")
|
||||
|
||||
# Collection for full articles
|
||||
articles_collection = client.get_or_create_collection("articles")
|
||||
|
||||
# Collection for company facts
|
||||
company_collection = client.get_or_create_collection("company_facts")
|
||||
|
||||
return facts_collection, articles_collection, company_collection
|
||||
|
||||
def embed_and_store_facts(facts, facts_collection, articles_collection, company_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')
|
||||
}]
|
||||
)
|
||||
|
||||
# Store company-specific facts in company collection
|
||||
if facts.get('key_entities', []):
|
||||
for entity in facts['key_entities']:
|
||||
# Check if it's a company (simplified - in reality you'd have a company detection system)
|
||||
company_fact_text = f"Company: {entity}. Key points: {', '.join(facts.get('main_points', []))}"
|
||||
company_embedding = get_embedding(company_fact_text)
|
||||
if company_embedding:
|
||||
company_collection.upsert(
|
||||
ids=[str(uuid.uuid4())],
|
||||
documents=[company_fact_text],
|
||||
embeddings=[company_embedding],
|
||||
metadatas=[{
|
||||
"company": entity,
|
||||
"source": facts['source'],
|
||||
"published": facts['published'],
|
||||
"filename": facts['filename'],
|
||||
"type": "company_fact",
|
||||
"processed_at": facts['processed_at'],
|
||||
"main_topic": facts.get('main_topic', 'Unknown')
|
||||
}]
|
||||
)
|
||||
|
||||
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, company_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,
|
||||
company_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()
|
||||
188
embedding/health_check.py
Executable file
188
embedding/health_check.py
Executable file
@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Health check script for the embedding pipeline with Prometheus metrics
|
||||
This script verifies that all components of the embedding pipeline are working correctly
|
||||
and exports metrics for Grafana visualization
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import requests
|
||||
import logging
|
||||
from datetime import datetime
|
||||
import sys
|
||||
from prometheus_client import start_http_server, Gauge, Counter, Histogram
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Prometheus metrics
|
||||
# Health status metrics
|
||||
health_status = Gauge('embedding_pipeline_health_status', 'Health status of pipeline components (1=healthy, 0=unhealthy)', ['component'])
|
||||
# Processing metrics
|
||||
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')
|
||||
# Database metrics
|
||||
chromadb_collections_count = Gauge('embedding_pipeline_chromadb_collections', 'Number of collections in ChromaDB')
|
||||
# Cache metrics
|
||||
cache_entries_count = Gauge('embedding_pipeline_cache_entries', 'Number of entries in cache file')
|
||||
|
||||
def check_chromadb_connection():
|
||||
"""Check if ChromaDB is accessible"""
|
||||
try:
|
||||
import chromadb
|
||||
CHROMADB_HOST = os.getenv("CHROMADB_HOST", "example.com")
|
||||
CHROMADB_PORT = int(os.getenv("CHROMADB_PORT", "8000"))
|
||||
|
||||
client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
|
||||
# Test connection by getting all collections
|
||||
collections = client.list_collections()
|
||||
chromadb_collections_count.set(len(collections))
|
||||
logger.info(f"✓ ChromaDB connection successful. Found {len(collections)} collections")
|
||||
health_status.labels(component='chromadb').set(1)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"✗ ChromaDB connection failed: {e}")
|
||||
health_status.labels(component='chromadb').set(0)
|
||||
return False
|
||||
|
||||
def check_ai_server_connection():
|
||||
"""Check if AI server is accessible"""
|
||||
try:
|
||||
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"
|
||||
|
||||
# Test by sending a simple request
|
||||
start_time = datetime.now()
|
||||
response = requests.post(
|
||||
AI_SERVER_URL,
|
||||
json={
|
||||
"input": "test",
|
||||
"model": "text-embedding-3-small"
|
||||
},
|
||||
timeout=10
|
||||
)
|
||||
response.raise_for_status()
|
||||
end_time = datetime.now()
|
||||
response_time = (end_time - start_time).total_seconds()
|
||||
logger.info(f"✓ AI server connection successful. Response time: {response_time:.2f}s")
|
||||
health_status.labels(component='ai_server').set(1)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"✗ AI server connection failed: {e}")
|
||||
health_status.labels(component='ai_server').set(0)
|
||||
return False
|
||||
|
||||
def check_scraper_directory():
|
||||
"""Check if scraper directory exists and has articles"""
|
||||
try:
|
||||
scraper_articles_dir = "/scraper/articles"
|
||||
if not os.path.exists(scraper_articles_dir):
|
||||
logger.error(f"✗ Scraper directory does not exist: {scraper_articles_dir}")
|
||||
health_status.labels(component='scraper_dir').set(0)
|
||||
return False
|
||||
|
||||
# Check if there are any JSON files
|
||||
json_files = []
|
||||
for root, dirs, files in os.walk(scraper_articles_dir):
|
||||
for file in files:
|
||||
if file.endswith('.json'):
|
||||
json_files.append(os.path.join(root, file))
|
||||
|
||||
if json_files:
|
||||
logger.info(f"✓ Scraper directory accessible. Found {len(json_files)} article files")
|
||||
health_status.labels(component='scraper_dir').set(1)
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"⚠ Scraper directory exists but no JSON files found")
|
||||
health_status.labels(component='scraper_dir').set(1) # Directory exists, just no files yet
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"✗ Error checking scraper directory: {e}")
|
||||
health_status.labels(component='scraper_dir').set(0)
|
||||
return False
|
||||
|
||||
def check_cache_file():
|
||||
"""Check if cache file is accessible"""
|
||||
try:
|
||||
CACHE_FILE = os.getenv("CACHE_FILE", "processed_articles_cache.json")
|
||||
if os.path.exists(CACHE_FILE):
|
||||
logger.info("✓ Cache file exists")
|
||||
# Try to read it
|
||||
with open(CACHE_FILE, 'r') as f:
|
||||
cache_data = json.load(f)
|
||||
cache_entries_count.set(len(cache_data))
|
||||
logger.info(f"✓ Cache file readable. Contains {len(cache_data)} entries")
|
||||
return True
|
||||
else:
|
||||
logger.info("✓ Cache file does not exist (this is normal for first run)")
|
||||
cache_entries_count.set(0)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"✗ Error reading cache file: {e}")
|
||||
return False
|
||||
|
||||
def collect_pipeline_metrics():
|
||||
"""Collect and export all pipeline metrics"""
|
||||
try:
|
||||
# This function would be called by the main pipeline to collect metrics
|
||||
logger.info("Collecting pipeline metrics...")
|
||||
# In a real implementation, this would be called by the main processing loop
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error collecting pipeline metrics: {e}")
|
||||
return False
|
||||
|
||||
def main():
|
||||
"""Run all health checks and start Prometheus server"""
|
||||
# Start Prometheus metrics server on port 8001
|
||||
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}")
|
||||
|
||||
logger.info("Starting embedding pipeline health check...")
|
||||
|
||||
checks = [
|
||||
check_chromadb_connection,
|
||||
check_ai_server_connection,
|
||||
check_scraper_directory,
|
||||
check_cache_file
|
||||
]
|
||||
|
||||
results = []
|
||||
for check in checks:
|
||||
try:
|
||||
result = check()
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
logger.error(f"Error running {check.__name__}: {e}")
|
||||
results.append(False)
|
||||
|
||||
# Summary
|
||||
passed = sum(results)
|
||||
total = len(results)
|
||||
|
||||
logger.info(f"\nHealth Check Summary: {passed}/{total} checks passed")
|
||||
|
||||
# Set overall health status
|
||||
overall_health = 1 if passed == total else 0
|
||||
health_status.labels(component='overall').set(overall_health)
|
||||
|
||||
if passed == total:
|
||||
logger.info("✓ All health checks passed - pipeline is ready to run")
|
||||
return 0
|
||||
else:
|
||||
logger.error("✗ Some health checks failed - pipeline may not work correctly")
|
||||
return 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
31
embedding/run_embedding_pipeline.sh
Executable file
31
embedding/run_embedding_pipeline.sh
Executable file
@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Embedding pipeline cron job script
|
||||
# This script runs the advanced embedding pipeline periodically
|
||||
|
||||
# Set working directory
|
||||
cd /embedding
|
||||
|
||||
# Create log directory if it doesn't exist
|
||||
mkdir -p logs
|
||||
|
||||
# Get current timestamp
|
||||
TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
|
||||
|
||||
# Log file
|
||||
LOG_FILE="logs/embedding_pipeline_$TIMESTAMP.log"
|
||||
|
||||
echo "Starting embedding pipeline at $(date)" >> $LOG_FILE
|
||||
|
||||
# Run the embedding pipeline
|
||||
python3 advanced_embedder.py >> $LOG_FILE 2>&1
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "Embedding pipeline completed successfully at $(date)" >> $LOG_FILE
|
||||
echo "SUCCESS: Embedding pipeline completed"
|
||||
else
|
||||
echo "Embedding pipeline failed at $(date)" >> $LOG_FILE
|
||||
echo "ERROR: Embedding pipeline failed"
|
||||
fi
|
||||
|
||||
echo "Embedding pipeline finished at $(date)" >> $LOG_FILE
|
||||
Loading…
x
Reference in New Issue
Block a user