139 lines
4.6 KiB
Python
Executable File
139 lines
4.6 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
|
|
|
|
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
|
|
client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
|
|
|
|
# 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() |