Spaces:
Runtime error
Runtime error
File size: 2,901 Bytes
c319b15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """CLI entry point for testing RAG functionality."""
from models import create_embeddings, create_llm
from qa_chain import create_qa_chain
from vectorstore import load_or_create_vectorstore
def main():
"""Main execution function for CLI testing."""
embeddings = create_embeddings()
vectorstore = load_or_create_vectorstore(embeddings)
qa_chain = create_qa_chain(vectorstore)
query = "Tell me about Arcee Fusion."
chat_history = []
# Test vanilla response
print("\n=== Vanilla Response (No RAG) ===")
streaming_llm = create_llm(streaming=True)
print("Answer: ", end="", flush=True)
try:
for chunk in streaming_llm.stream(query):
print(chunk.content, end="", flush=True)
print()
except Exception as e:
print(f"\nError: {e}")
# Test RAG response with MMR (diverse results)
print("\n=== RAG Response (MMR for diversity) ===")
print("Answer: ", end="", flush=True)
chunk_data = None
try:
for chunk_data in qa_chain.stream(
{"question": query, "chat_history": chat_history}
):
print(chunk_data["chunk"], end="", flush=True)
print()
# Print sources
if chunk_data and chunk_data.get("source_documents"):
print("\nSources:")
seen_sources = set()
for doc in chunk_data["source_documents"]:
source = doc.metadata.get("source", "Unknown")
page = doc.metadata.get("page", "unknown")
source_key = f"{source}:{page}"
if source_key not in seen_sources:
print(f"- {source}, page {page}")
seen_sources.add(source_key)
except Exception as e:
print(f"\nError: {e}")
# Example: RAG with metadata filter (filter by page number)
print("\n=== RAG Response with Metadata Filter (page >= 5) ===")
print("Answer: ", end="", flush=True)
chunk_data = None
try:
for chunk_data in qa_chain.stream(
{
"question": query,
"chat_history": chat_history,
"filter": {"page": {"$gte": 5}}, # Only pages 5+
}
):
print(chunk_data["chunk"], end="", flush=True)
print()
if chunk_data and chunk_data.get("source_documents"):
print("\nSources (filtered):")
seen_sources = set()
for doc in chunk_data["source_documents"]:
source = doc.metadata.get("source", "Unknown")
page = doc.metadata.get("page", "unknown")
source_key = f"{source}:{page}"
if source_key not in seen_sources:
print(f"- {source}, page {page}")
seen_sources.add(source_key)
except Exception as e:
print(f"\nError: {e}")
if __name__ == "__main__":
main()
|