StockDocs/embedding/query_chromadb.py

#!/usr/bin/env python3
"""
Script to demonstrate querying ChromaDB instance
This shows how to connect and query your ChromaDB database
"""

import os
import chromadb
from chromadb.config import Settings

def connect_to_chromadb():
    """Connect to ChromaDB instance"""
    try:
        # Get connection details from environment or use defaults
        CHROMADB_HOST = os.getenv("CHROMADB_HOST", "example.com")
        CHROMADB_PORT = int(os.getenv("CHROMADB_PORT", "8000"))

        print(f"Connecting to ChromaDB at {CHROMADB_HOST}:{CHROMADB_PORT}")

        # Create client connection using Settings for remote server
        settings = Settings(
            chroma_api_impl="rest",
            chroma_server_host=CHROMADB_HOST,
            chroma_server_http_port=CHROMADB_PORT,
            chroma_server_ssl_enabled=False
        )

        print(f"Creating client with settings: {settings}")
        client = chromadb.Client(settings=settings)

        # Test connection
        collections = client.list_collections()
        print(f"Successfully connected! Found {len(collections)} collections:")
        for collection in collections:
            print(f"  - {collection.name}")

        return client

    except Exception as e:
        print(f"Failed to connect to ChromaDB: {e}")
        return None

def query_facts_collection(client):
    """Query the facts collection"""
    try:
        # Get the facts collection
        facts_collection = client.get_collection("facts")

        print("\n=== Querying Facts Collection ===")

        # Example queries that would work with your data
        example_queries = [
            "Unrivaled attendance records",
            "Unrivaled league revenue",
            "David Levy on Unrivaled",
            "Fox Business coverage of Unrivaled"
        ]

        for i, query in enumerate(example_queries, 1):
            print(f"\n{i}. Query: '{query}'")

            # Perform similarity search
            results = facts_collection.query(
                query_texts=[query],
                n_results=2,
                include=["documents", "metadatas", "distances"]
            )

            if results['documents'] and len(results['documents'][0]) > 0:
                print("   Results:")
                for j, doc in enumerate(results['documents'][0], 1):
                    print(f"     {j}. {doc[:100]}...")
                    if j >= 2:  # Show only first 2 results
                        break
            else:
                print("   No results found")

    except Exception as e:
        print(f"Error querying facts collection: {e}")

def query_articles_collection(client):
    """Query the articles collection"""
    try:
        # Get the articles collection
        articles_collection = client.get_collection("articles")

        print("\n=== Querying Articles Collection ===")

        # Example query
        query = "Unrivaled women's basketball"
        print(f"Query: '{query}'")

        # Perform similarity search
        results = articles_collection.query(
            query_texts=[query],
            n_results=2,
            include=["documents", "metadatas", "distances"]
        )

        if results['documents'] and len(results['documents'][0]) > 0:
            print("Results:")
            for j, doc in enumerate(results['documents'][0], 1):
                print(f"  {j}. {doc[:100]}...")
        else:
            print("No results found")

    except Exception as e:
        print(f"Error querying articles collection: {e}")

def main():
    """Main function to demonstrate ChromaDB queries"""
    print("=== ChromaDB Query Demonstration ===")
    print("This script shows how to connect and query your ChromaDB instance")
    print()

    # Connect to ChromaDB
    client = connect_to_chromadb()

    if client:
        print("\n=== Available Collections ===")
        collections = client.list_collections()
        for collection in collections:
            print(f"  - {collection.name} ({collection.count()} items)")

        # Query each collection
        query_facts_collection(client)
        query_articles_collection(client)

        print("\n=== Query Capabilities ===")
        print("The system supports:")
        print("✓ Semantic similarity search across facts")
        print("✓ Entity-based filtering")
        print("✓ Source-based filtering")
        print("✓ Time-based filtering")
        print("✓ Multi-entity queries")
        print()
        print("All queries use the qwen3:8b embedding model for efficient searching")

    else:
        print("Cannot connect to ChromaDB. Please ensure:")
        print("1. ChromaDB server is running")
        print("2. Network connectivity to the server")
        print("3. Correct host/port configuration")

if __name__ == "__main__":
    main()