Initial commit: StockDocs data API server
Carved out from the StockDocs monorepo. Flask API over the ChromaDB vector store and article corpus with OpenAPI spec.
This commit is contained in:
commit
e977b052b9
20
Dockerfile
Normal file
20
Dockerfile
Normal file
@ -0,0 +1,20 @@
|
|||||||
|
# Use an official Python runtime as a base image
|
||||||
|
FROM python:3.12.3-slim
|
||||||
|
|
||||||
|
# Set the working directory in the container
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
# Install Dependencies
|
||||||
|
COPY requirements.txt /app/
|
||||||
|
|
||||||
|
RUN pip install --no-cache-dir -r requirements.txt
|
||||||
|
|
||||||
|
# Copy the current directory contents into the container at /app
|
||||||
|
COPY . /app
|
||||||
|
|
||||||
|
# Expose the port the app runs on
|
||||||
|
EXPOSE 5005
|
||||||
|
|
||||||
|
# Run the Flask app
|
||||||
|
CMD ["python", "server.py"]
|
||||||
|
|
||||||
21
LICENSE
Normal file
21
LICENSE
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2025 Jarian Cottingham
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
||||||
51
README.md
Normal file
51
README.md
Normal file
@ -0,0 +1,51 @@
|
|||||||
|
# StockDocs — MCP
|
||||||
|
|
||||||
|
Data API server for the StockDocs financial news platform: exposes the processed article corpus and vector database over HTTP so LLM clients (via MCP wrappers) can query semantic search results, article content, and company facts.
|
||||||
|
|
||||||
|
Part of the StockDocs project family:
|
||||||
|
|
||||||
|
| Repo | What it is |
|
||||||
|
|------|------------|
|
||||||
|
| [StockDocs](https://git.jarianc.com/jarianc/StockDocs) | Processing core — article server, NLP analysis, embeddings |
|
||||||
|
| [stockdocs-scraper](https://git.jarianc.com/jarianc/stockdocs-scraper) | RSS scraper — collects the article corpus |
|
||||||
|
|
||||||
|
## Endpoints
|
||||||
|
|
||||||
|
| Method | Path | Description |
|
||||||
|
|--------|------|-------------|
|
||||||
|
| POST | `/query` | Semantic search over the ChromaDB vector store |
|
||||||
|
| POST | `/articles/query` | Filter the article corpus by criteria |
|
||||||
|
| GET | `/articles/latest/<field>` | Latest articles by field |
|
||||||
|
| GET | `/company/<name>/facts` | Company fact sheet |
|
||||||
|
| GET | `/company/<name>/products` | Company product list |
|
||||||
|
| POST | `/company/facts/update` | Update a company fact sheet |
|
||||||
|
| POST | `/facts` | Extracted-facts query |
|
||||||
|
| GET | `/health` | Health check |
|
||||||
|
| GET | `/info` | Service info |
|
||||||
|
| GET | `/openapi.json` | OpenAPI 3.0 spec |
|
||||||
|
|
||||||
|
## Running
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install -r requirements.txt
|
||||||
|
python server.py # :5005
|
||||||
|
```
|
||||||
|
|
||||||
|
Requires a running ChromaDB instance (host/port via environment).
|
||||||
|
|
||||||
|
## Docker
|
||||||
|
|
||||||
|
```bash
|
||||||
|
docker build -t stockdocs-mcp .
|
||||||
|
docker run -p 5005:5005 -e CHROMADB_HOST=<host> stockdocs-mcp
|
||||||
|
```
|
||||||
|
|
||||||
|
## Project Structure
|
||||||
|
|
||||||
|
```
|
||||||
|
.
|
||||||
|
├── server.py # Flask app — all endpoints
|
||||||
|
├── openapi.json # OpenAPI 3.0 spec (also served at /openapi.json)
|
||||||
|
├── requirements.txt # Python dependencies
|
||||||
|
└── Dockerfile # python:3.12-slim image
|
||||||
|
```
|
||||||
316
openapi.json
Normal file
316
openapi.json
Normal file
@ -0,0 +1,316 @@
|
|||||||
|
{
|
||||||
|
"openapi": "3.0.0",
|
||||||
|
"info": {
|
||||||
|
"title": "StockDocs",
|
||||||
|
"version": "1.0.0"
|
||||||
|
},
|
||||||
|
"paths": {
|
||||||
|
"/query": {
|
||||||
|
"post": {
|
||||||
|
"summary": "Query the vector database",
|
||||||
|
"requestBody": {
|
||||||
|
"required": true,
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"question": { "type": "string" }
|
||||||
|
},
|
||||||
|
"required": ["question"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Query results",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"results": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"document": { "type": "string" },
|
||||||
|
"score": { "type": "number" },
|
||||||
|
"metadata": { "type": "object" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/articles/query": {
|
||||||
|
"post": {
|
||||||
|
"summary": "Query articles based on a question with diversity filtering",
|
||||||
|
"requestBody": {
|
||||||
|
"required": true,
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"question": { "type": "string" },
|
||||||
|
"max_results": { "type": "integer" }
|
||||||
|
},
|
||||||
|
"required": ["question"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Query results with diverse articles",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"results": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"document": { "type": "string" },
|
||||||
|
"score": { "type": "number" },
|
||||||
|
"metadata": { "type": "object" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"query": { "type": "string" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/articles/latest/{field}": {
|
||||||
|
"get": {
|
||||||
|
"summary": "Get latest articles about a specific field with diversity",
|
||||||
|
"parameters": [
|
||||||
|
{
|
||||||
|
"name": "field",
|
||||||
|
"in": "path",
|
||||||
|
"required": true,
|
||||||
|
"schema": {
|
||||||
|
"type": "string"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Latest diverse articles for the field",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"results": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"document": { "type": "string" },
|
||||||
|
"score": { "type": "number" },
|
||||||
|
"metadata": { "type": "object" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"field": { "type": "string" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/company/{company_name}/facts": {
|
||||||
|
"get": {
|
||||||
|
"summary": "Get facts about a specific company",
|
||||||
|
"parameters": [
|
||||||
|
{
|
||||||
|
"name": "company_name",
|
||||||
|
"in": "path",
|
||||||
|
"required": true,
|
||||||
|
"schema": {
|
||||||
|
"type": "string"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Company facts",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"company": { "type": "string" },
|
||||||
|
"facts": { "type": "object" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/company/{company_name}/products": {
|
||||||
|
"get": {
|
||||||
|
"summary": "Get products information for a company",
|
||||||
|
"parameters": [
|
||||||
|
{
|
||||||
|
"name": "company_name",
|
||||||
|
"in": "path",
|
||||||
|
"required": true,
|
||||||
|
"schema": {
|
||||||
|
"type": "string"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Company products",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"company": { "type": "string" },
|
||||||
|
"products": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"name": { "type": "string" },
|
||||||
|
"release_date": { "type": "string" },
|
||||||
|
"specifications": { "type": "string" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/company/facts/update": {
|
||||||
|
"post": {
|
||||||
|
"summary": "Update or add company facts",
|
||||||
|
"requestBody": {
|
||||||
|
"required": true,
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"company_name": { "type": "string" },
|
||||||
|
"facts": { "type": "object" }
|
||||||
|
},
|
||||||
|
"required": ["company_name", "facts"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Facts updated successfully",
|
||||||
|
"content": {
|
||||||
|
"application/json": {
|
||||||
|
"schema": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"message": { "type": "string" },
|
||||||
|
"company": { "type": "string" },
|
||||||
|
"facts": { "type": "object" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/health": {
|
||||||
|
"get": {
|
||||||
|
"summary": "Health check",
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "OK"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/info": {
|
||||||
|
"get": {
|
||||||
|
"summary": "Service info",
|
||||||
|
"responses": {
|
||||||
|
"200": {
|
||||||
|
"description": "Info"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"/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" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
6
requirements.txt
Normal file
6
requirements.txt
Normal file
@ -0,0 +1,6 @@
|
|||||||
|
Flask
|
||||||
|
requests
|
||||||
|
numpy
|
||||||
|
pandas
|
||||||
|
scikit-learn
|
||||||
|
chromadb
|
||||||
0
requirements_clean.txt
Normal file
0
requirements_clean.txt
Normal file
412
server.py
Normal file
412
server.py
Normal file
@ -0,0 +1,412 @@
|
|||||||
|
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)
|
||||||
Loading…
x
Reference in New Issue
Block a user