297 lines
11 KiB
Python
297 lines
11 KiB
Python
import chromadb
|
|
from flask import Flask, request, jsonify, send_from_directory
|
|
import os
|
|
import json
|
|
import requests
|
|
import logging
|
|
from datetime import datetime
|
|
|
|
# Setup logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = Flask(__name__)
|
|
|
|
# ChromaDB client setup
|
|
try:
|
|
client = chromadb.HttpClient(host="chromadb", port=8000)
|
|
logger.info("Connected to ChromaDB successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to ChromaDB: {e}")
|
|
client = None
|
|
|
|
# Company facts database (in-memory for now, could be replaced with persistent storage)
|
|
company_facts = {
|
|
"Apple": {
|
|
"products": [
|
|
{"name": "iPhone", "release_date": "2007", "specifications": "Smartphone with iOS"},
|
|
{"name": "MacBook", "release_date": "2006", "specifications": "Laptop with M1 chip"},
|
|
{"name": "iPad", "release_date": "2010", "specifications": "Tablet with iOS"},
|
|
{"name": "Apple Watch", "release_date": "2015", "specifications": "Smartwatch with watchOS"}
|
|
],
|
|
"founded": "1976",
|
|
"ceo": "Tim Cook",
|
|
"headquarters": "Cupertino, California",
|
|
"ipo_year": "1980",
|
|
"market_cap": "$2.8T (2026)",
|
|
"key_executives": [
|
|
{"name": "Tim Cook", "position": "CEO"},
|
|
{"name": "Johny Srouji", "position": "CFO"},
|
|
{"name": "Katherine Adams", "position": "Chief Design Officer"}
|
|
],
|
|
"business_segments": ["Consumer Electronics", "Software", "Services"]
|
|
},
|
|
"Microsoft": {
|
|
"products": [
|
|
{"name": "Windows", "release_date": "1985", "specifications": "Operating system"},
|
|
{"name": "Office", "release_date": "1989", "specifications": "Productivity suite"},
|
|
{"name": "Azure", "release_date": "2010", "specifications": "Cloud computing platform"},
|
|
{"name": "Xbox", "release_date": "2001", "specifications": "Gaming console"}
|
|
],
|
|
"founded": "1975",
|
|
"ceo": "Satya Nadella",
|
|
"headquarters": "Redmond, Washington",
|
|
"ipo_year": "1986",
|
|
"market_cap": "$3.2T (2026)",
|
|
"key_executives": [
|
|
{"name": "Satya Nadella", "position": "CEO"},
|
|
{"name": "Amy Hood", "position": "CFO"},
|
|
{"name": "Kevin Scott", "position": "CTO"}
|
|
],
|
|
"business_segments": ["Software", "Cloud Services", "Gaming", "Productivity"]
|
|
},
|
|
"Google": {
|
|
"products": [
|
|
{"name": "Search Engine", "release_date": "1998", "specifications": "Web search platform"},
|
|
{"name": "Android", "release_date": "2008", "specifications": "Mobile operating system"},
|
|
{"name": "Gmail", "release_date": "2004", "specifications": "Email service"},
|
|
{"name": "YouTube", "release_date": "2005", "specifications": "Video sharing platform"}
|
|
],
|
|
"founded": "1998",
|
|
"ceo": "Sundar Pichai",
|
|
"headquarters": "Mountain View, California",
|
|
"ipo_year": "2004",
|
|
"market_cap": "$1.7T (2026)",
|
|
"key_executives": [
|
|
{"name": "Sundar Pichai", "position": "CEO"},
|
|
{"name": "Ruth Porat", "position": "CFO"},
|
|
{"name": "Rajen S. Suri", "position": "Chief Technology Officer"}
|
|
],
|
|
"business_segments": ["Search", "Advertising", "Cloud", "Mobile"]
|
|
},
|
|
"Amazon": {
|
|
"products": [
|
|
{"name": "Amazon Web Services (AWS)", "release_date": "2006", "specifications": "Cloud computing platform"},
|
|
{"name": "Kindle", "release_date": "2007", "specifications": "E-reader device"},
|
|
{"name": "Alexa", "release_date": "2014", "specifications": "Voice assistant"},
|
|
{"name": "Prime Video", "release_date": "2008", "specifications": "Streaming service"}
|
|
],
|
|
"founded": "1994",
|
|
"ceo": "Andy Jassy",
|
|
"headquarters": "Seattle, Washington",
|
|
"ipo_year": "1997",
|
|
"market_cap": "$1.5T (2026)",
|
|
"key_executives": [
|
|
{"name": "Andy Jassy", "position": "CEO"},
|
|
{"name": "Brian T. Olsavsky", "position": "CFO"},
|
|
{"name": "Wendy J. Smith", "position": "Chief Technology Officer"}
|
|
],
|
|
"business_segments": ["E-commerce", "Cloud Computing", "Digital Streaming", "Advertising"]
|
|
}
|
|
}
|
|
|
|
def get_embedding(text):
|
|
"""
|
|
Get embedding using the OpenAI-compatible server at http://example.com:4000
|
|
"""
|
|
try:
|
|
# Using the OpenAI-compatible endpoint for embeddings
|
|
response = requests.post(
|
|
"http://example.com:4000/v1/embeddings",
|
|
json={
|
|
"input": text,
|
|
"model": "text-embedding-3-small" # or whatever model you're using
|
|
},
|
|
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 get_diverse_articles(articles, max_diverse=5):
|
|
"""
|
|
Filter articles to ensure diversity in content, sources, and perspectives
|
|
"""
|
|
if len(articles) <= max_diverse:
|
|
return articles
|
|
|
|
# More sophisticated diversity algorithm
|
|
diverse_articles = []
|
|
source_count = {}
|
|
topic_count = {}
|
|
|
|
# First pass: try to get articles from different sources
|
|
for article in articles:
|
|
source = article.get('metadata', {}).get('source', 'unknown')
|
|
topic = article.get('metadata', {}).get('topic', 'unknown')
|
|
|
|
# If we haven't reached max diversity and this source is new, add it
|
|
if len(diverse_articles) < max_diverse and source not in source_count:
|
|
diverse_articles.append(article)
|
|
source_count[source] = 1
|
|
topic_count[topic] = topic_count.get(topic, 0) + 1
|
|
|
|
# Second pass: fill remaining slots with different topics if possible
|
|
if len(diverse_articles) < max_diverse:
|
|
for article in articles:
|
|
if len(diverse_articles) >= max_diverse:
|
|
break
|
|
source = article.get('metadata', {}).get('source', 'unknown')
|
|
topic = article.get('metadata', {}).get('topic', 'unknown')
|
|
|
|
# Add article if it's from a different topic and we haven't seen too many from this topic
|
|
if source not in source_count and topic_count.get(topic, 0) < 2:
|
|
diverse_articles.append(article)
|
|
source_count[source] = 1
|
|
topic_count[topic] = topic_count.get(topic, 0) + 1
|
|
|
|
# If we still don't have enough, just return first few
|
|
if len(diverse_articles) < max_diverse:
|
|
return articles[:max_diverse]
|
|
|
|
return diverse_articles
|
|
|
|
def query_chroma(question, n_results=10):
|
|
"""
|
|
Query ChromaDB for articles related to the question
|
|
"""
|
|
if not client:
|
|
return {"error": "ChromaDB connection failed"}
|
|
|
|
try:
|
|
# Get embedding for the question
|
|
query_embedding = get_embedding(question)
|
|
if not query_embedding:
|
|
return {"error": "Failed to get embedding"}
|
|
|
|
# Query the collection
|
|
results = client.get_or_create_collection("news").query(
|
|
query_embeddings=[query_embedding],
|
|
n_results=n_results,
|
|
)
|
|
|
|
return results
|
|
except Exception as e:
|
|
logger.error(f"Error querying ChromaDB: {e}")
|
|
return {"error": f"Query failed: {str(e)}"}
|
|
|
|
@app.route("/facts", methods=["POST"])
|
|
def get_facts():
|
|
"""
|
|
Get the best set of facts about a question
|
|
"""
|
|
data = request.get_json()
|
|
question = data.get("question")
|
|
|
|
if not question:
|
|
return jsonify({"error": "Missing 'question' in request body"}), 400
|
|
|
|
# Query ChromaDB for relevant articles
|
|
results = query_chroma(question, n_results=10)
|
|
|
|
if "error" in results:
|
|
return jsonify({"error": results["error"]}), 500
|
|
|
|
# Process results to create diverse article set
|
|
mcp_results = []
|
|
for doc, score, meta in zip(
|
|
results.get("documents", [[]])[0],
|
|
results.get("distances", [[]])[0],
|
|
results.get("metadatas", [[]])[0]):
|
|
mcp_results.append({
|
|
"document": doc,
|
|
"score": float(score),
|
|
"metadata": meta
|
|
})
|
|
|
|
# Apply diversity filtering
|
|
diverse_results = get_diverse_articles(mcp_results, 5)
|
|
|
|
# Combine with company facts if question mentions a company
|
|
company_facts_result = {}
|
|
question_lower = question.lower()
|
|
|
|
# Check if question mentions any known company
|
|
for company_name in company_facts.keys():
|
|
if company_name.lower() in question_lower:
|
|
company_facts_result = company_facts[company_name]
|
|
break
|
|
|
|
# Return combined results
|
|
response_data = {
|
|
"question": question,
|
|
"articles": diverse_results,
|
|
"company_facts": company_facts_result,
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
|
|
return jsonify(response_data)
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def health():
|
|
"""Health check endpoint"""
|
|
return jsonify({"status": "healthy", "timestamp": datetime.now().isoformat()})
|
|
|
|
@app.route("/info", methods=["GET"])
|
|
def info():
|
|
"""Service information endpoint"""
|
|
return jsonify({
|
|
"provider": "StockDoc",
|
|
"service": "MCP Server",
|
|
"version": "1.0.0",
|
|
"collection": "news",
|
|
"embedding_service": "http://example.com:4000"
|
|
})
|
|
|
|
@app.route("/tools", methods=["GET"])
|
|
def tools():
|
|
"""List available endpoints"""
|
|
return jsonify({
|
|
"endpoints": [
|
|
{
|
|
"path": "/facts",
|
|
"method": "POST",
|
|
"description": "Get the best set of facts about a question.",
|
|
"request_format": {"question": "string"},
|
|
"response_format": {
|
|
"question": "string",
|
|
"articles": [
|
|
{"document": "string", "score": "float", "metadata": "object"}
|
|
],
|
|
"company_facts": "object",
|
|
"timestamp": "string"
|
|
}
|
|
},
|
|
{
|
|
"path": "/health",
|
|
"method": "GET",
|
|
"description": "Health check endpoint."
|
|
},
|
|
{
|
|
"path": "/info",
|
|
"method": "GET",
|
|
"description": "Service information endpoint."
|
|
}
|
|
]
|
|
})
|
|
|
|
@app.route("/openapi.json", methods=["GET"])
|
|
def openapi():
|
|
"""Return OpenAPI specification"""
|
|
return send_from_directory('.', 'openapi.json')
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5005, threaded=True, debug=True) |