StockDocs/embedding/advanced_embedder.py

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 with better error handling
try:
    logging.basicConfig(
        level=logging.INFO, 
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler('embedding_pipeline.log'),
            logging.StreamHandler()
        ]
    )
except Exception as e:
    # Fallback if file logging fails
    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}")
    # In cron job environment, we might want to exit gracefully or continue with logging
    # For now, let's continue but log the error
    pass

# 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()