Jarian Cottingham
e977b052b9
Carved out from the StockDocs monorepo. Flask API over the ChromaDB vector store and article corpus with OpenAPI spec.
413 lines
14 KiB
Python
413 lines
14 KiB
Python
import chromadb
|
|
from flask import Flask, request, jsonify, send_from_directory
|
|
import os
|
|
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
|
|
headers = {
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
# Add API key if available
|
|
api_key = os.getenv("AI_SERVICE_API_KEY")
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
response = requests.post(
|
|
"http://example.com:4000/v1/embeddings",
|
|
json={
|
|
"input": text,
|
|
"model": "text-embedding-3-small" # or whatever model you're using
|
|
},
|
|
headers=headers,
|
|
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)}"}
|
|
|
|
# New endpoints implementation based on the OpenAPI specification
|
|
|
|
@app.route("/query", methods=["POST"])
|
|
def query_vector_database():
|
|
"""Query the vector database"""
|
|
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)
|
|
|
|
return jsonify({"results": diverse_results})
|
|
|
|
@app.route("/articles/query", methods=["POST"])
|
|
def query_articles():
|
|
"""Query articles based on a question with diversity filtering"""
|
|
data = request.get_json()
|
|
question = data.get("question")
|
|
max_results = data.get("max_results", 5)
|
|
|
|
if not question:
|
|
return jsonify({"error": "Missing 'question' in request body"}), 400
|
|
|
|
# Query ChromaDB for relevant articles
|
|
results = query_chroma(question, n_results=max_results)
|
|
|
|
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, max_results)
|
|
|
|
return jsonify({
|
|
"results": diverse_results,
|
|
"query": question
|
|
})
|
|
|
|
@app.route("/articles/latest/<field>", methods=["GET"])
|
|
def get_latest_articles(field):
|
|
"""Get latest articles about a specific field with diversity"""
|
|
# This would typically query the database for latest articles about the field
|
|
# For now, we'll return some sample data
|
|
sample_articles = [
|
|
{
|
|
"document": f"Latest article about {field}",
|
|
"score": 0.95,
|
|
"metadata": {
|
|
"source": "Sample Source",
|
|
"topic": field,
|
|
"date": "2026-01-31"
|
|
}
|
|
}
|
|
]
|
|
|
|
return jsonify({
|
|
"results": sample_articles,
|
|
"field": field
|
|
})
|
|
|
|
@app.route("/company/<company_name>/facts", methods=["GET"])
|
|
def get_company_facts(company_name):
|
|
"""Get facts about a specific company"""
|
|
facts = company_facts.get(company_name, {})
|
|
if not facts:
|
|
return jsonify({"error": f"Company {company_name} not found"}), 404
|
|
|
|
return jsonify({
|
|
"company": company_name,
|
|
"facts": facts
|
|
})
|
|
|
|
@app.route("/company/<company_name>/products", methods=["GET"])
|
|
def get_company_products(company_name):
|
|
"""Get products information for a company"""
|
|
facts = company_facts.get(company_name, {})
|
|
products = facts.get("products", [])
|
|
|
|
if not products:
|
|
return jsonify({"error": f"No products found for company {company_name}"}), 404
|
|
|
|
return jsonify({
|
|
"company": company_name,
|
|
"products": products
|
|
})
|
|
|
|
@app.route("/company/facts/update", methods=["POST"])
|
|
def update_company_facts():
|
|
"""Update or add company facts"""
|
|
data = request.get_json()
|
|
company_name = data.get("company_name")
|
|
facts = data.get("facts")
|
|
|
|
if not company_name or not facts:
|
|
return jsonify({"error": "Missing 'company_name' or 'facts' in request body"}), 400
|
|
|
|
# Update or add company facts
|
|
company_facts[company_name] = facts
|
|
|
|
return jsonify({
|
|
"message": "Facts updated successfully",
|
|
"company": company_name,
|
|
"facts": facts
|
|
})
|
|
|
|
# Restore the original /facts endpoint
|
|
@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("/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)
|