Spaces:
Running
Running
File size: 1,803 Bytes
49cf970 | 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 | import os
import sys
from pathlib import Path
# Add project root to sys.path
root_dir = Path(__file__).parent.parent
sys.path.append(str(root_dir))
from dotenv import load_dotenv
from agent.agent import LlamaPDFAgent
import io
load_dotenv()
def test_agent():
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
print("GROQ_API_KEY not found in environment.")
return
agent = LlamaPDFAgent(api_key=api_key)
# Use the downloaded NVIDIA PDF - updated path
pdf_path = os.path.join(root_dir, "nvidia_q4_fy24.pdf")
if not os.path.exists(pdf_path):
print(f"PDF not found: {pdf_path}")
return
with open(pdf_path, "rb") as f:
# Mocking a streamlit-like upload object
class MockFile:
def __init__(self, file, name):
self.file = file
self.name = name
def read(self):
return self.file.read()
def seek(self, pos):
self.file.seek(pos)
def tell(self):
return self.file.tell()
mock_file = MockFile(f, pdf_path)
print("Ingesting PDF...")
msg = agent.ingest_pdf(mock_file)
print(msg)
print("\n--- Testing Q&A ---")
q = "What was the total revenue for FY24?"
result = agent.answer_question(q)
print(f"Q: {q}")
print(f"A: {result['answer']}")
print("\nSources:")
for src in result['sources']:
print(f"- [Page {src['page']}] {src['text'][:100]}...")
print("\n--- Testing Deep Insights ---")
insights = agent.get_deep_insights()
for key, value in insights.items():
print(f"\n[{key.upper()}]")
print(value)
if __name__ == "__main__":
test_agent()
|