148 lines
4.9 KiB
Python
Executable File
148 lines
4.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Script to demonstrate querying ChromaDB instance
|
|
This shows how to connect and query your ChromaDB database
|
|
"""
|
|
|
|
import os
|
|
import chromadb
|
|
import json
|
|
from pathlib import Path
|
|
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() |