import os import json import chromadb import uuid import datetime import requests import logging 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: # 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 - using the centralized endpoint 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}" # 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}") # API Key for AI service authentication AI_SERVICE_API_KEY = os.getenv("AI_SERVICE_API_KEY") # Batch processing configuration BATCH_SIZE = int(os.getenv("EMBEDDING_BATCH_SIZE", "50")) def get_embedding(text): """ Get embedding using the OpenAI-compatible server with qwen3:8b model """ try: embedding_url = f"{AI_SERVER_URL}/v1/embeddings" # Build headers with authentication if available headers = { "Content-Type": "application/json" } if AI_SERVICE_API_KEY: headers["Authorization"] = f"Bearer {AI_SERVICE_API_KEY}" response = requests.post( embedding_url, json={ "input": text, "model": "qwen3:8b" }, headers=headers, timeout=60 ) 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 the centralized AI service with gpt-oss model """ try: # Use the centralized AI endpoint for fact extraction extraction_url = f"{AI_SERVER_URL}/v1/chat/completions" # Build headers with authentication if available headers = { "Content-Type": "application/json" } if AI_SERVICE_API_KEY: headers["Authorization"] = f"Bearer {AI_SERVICE_API_KEY}" # 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[:2000]}... 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": "{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"] }} """ # Call the AI service with gpt-oss model for fact extraction response = requests.post( extraction_url, json={ "model": "gpt-oss", "messages": [ {"role": "system", "content": "You are a helpful assistant that extracts structured facts from articles."}, {"role": "user", "content": prompt} ], "temperature": 0.3, "max_tokens": 1000 }, headers=headers, timeout=60 ) response.raise_for_status() # Parse the response result = response.json() extracted_text = result['choices'][0]['message']['content'].strip() # Try to parse the JSON from the response try: facts = json.loads(extracted_text) except json.JSONDecodeError: # If JSON parsing fails, create a basic structure 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 with batch processing """ 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() # Collect all articles to process articles_to_process = [] 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 not in processed_cache: articles_to_process.append((file_path, file)) logger.info(f"Found {len(articles_to_process)} articles to process in batches of {BATCH_SIZE}") # Process articles in batches total_processed = 0 total_failed = 0 for i in range(0, len(articles_to_process), BATCH_SIZE): batch = articles_to_process[i:i + BATCH_SIZE] logger.info(f"Processing batch {i//BATCH_SIZE + 1} with {len(batch)} articles") batch_processed = 0 batch_failed = 0 for file_path, file in batch: try: # 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: # Update cache with detailed status tracking processed_cache[file_path] = { "processed_date": datetime.datetime.now().isoformat(), "status": "fact_extracted", "embedding_status": "pending", "last_updated": datetime.datetime.now().isoformat() } articles_processed_total.inc() logger.info(f"Successfully processed {file}") batch_processed += 1 total_processed += 1 else: articles_failed_total.inc() logger.error(f"Failed to embed {file}") batch_failed += 1 total_failed += 1 else: logger.error(f"Failed to extract facts for {file}") batch_failed += 1 total_failed += 1 except Exception as e: logger.error(f"Error processing article {file}: {e}") articles_failed_total.inc() batch_failed += 1 total_failed += 1 logger.info(f"Batch {i//BATCH_SIZE + 1} completed: {batch_processed} successful, {batch_failed} failed") # Save cache periodically during batch processing try: with open(CACHE_FILE, 'w', encoding='utf-8') as f: json.dump(processed_cache, f, indent=2) logger.info(f"Cache updated after batch {i//BATCH_SIZE + 1}") except Exception as e: logger.error(f"Error saving cache file: {e}") # Final cache save try: with open(CACHE_FILE, 'w', encoding='utf-8') as f: json.dump(processed_cache, f, indent=2) logger.info(f"Final cache saved with {len(processed_cache)} entries") except Exception as e: logger.error(f"Error saving final 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(f"Total processed: {total_processed}, Total failed: {total_failed}") logger.info("Embedding pipeline completed") if __name__ == "__main__": main()