mirror of
https://github.com/MODSetter/SurfSense.git
synced 2025-09-03 19:19:10 +00:00
126 lines
No EOL
4.6 KiB
Python
126 lines
No EOL
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for the Researcher LangGraph agent.
|
|
|
|
This script demonstrates how to invoke the researcher agent with a sample query.
|
|
Run this script directly from VSCode using the "Run Python File" button or
|
|
right-click and select "Run Python File in Terminal".
|
|
|
|
Before running:
|
|
1. Make sure your Python environment has all required dependencies
|
|
2. Create a .env file with any required API keys
|
|
3. Ensure database connection is properly configured
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add project root to Python path so that 'app' can be found as a module
|
|
# Get the absolute path to the surfsense_backend directory which contains the app module
|
|
project_root = str(Path(__file__).resolve().parents[3]) # Go up 3 levels from the script: app/agents/researcher -> app/agents -> app -> surfsense_backend
|
|
print(f"Adding to Python path: {project_root}")
|
|
sys.path.insert(0, project_root)
|
|
|
|
# Now import the modules after fixing the path
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
|
|
from sqlalchemy.orm import sessionmaker
|
|
from dotenv import load_dotenv
|
|
|
|
# These imports should now work with the correct path
|
|
from app.agents.researcher.graph import graph
|
|
from app.agents.researcher.state import State
|
|
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Database connection string - use a test database or mock
|
|
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost:5432/surfsense"
|
|
|
|
# Create async engine and session
|
|
engine = create_async_engine(DATABASE_URL)
|
|
async_session_maker = sessionmaker(
|
|
engine, class_=AsyncSession, expire_on_commit=False
|
|
)
|
|
|
|
async def run_test():
|
|
"""Run a test of the researcher agent."""
|
|
print("Starting researcher agent test...")
|
|
|
|
# Create a database session
|
|
async with async_session_maker() as db_session:
|
|
# Sample configuration
|
|
config = {
|
|
"configurable": {
|
|
"user_query": "What are the best clash royale decks recommended by Surgical Goblin?",
|
|
"num_sections": 1,
|
|
"connectors_to_search": [
|
|
"YOUTUBE_VIDEO",
|
|
],
|
|
"user_id": "d6ac2187-7407-4664-8734-af09926d161e",
|
|
"search_space_id": 2
|
|
}
|
|
}
|
|
|
|
try:
|
|
# Initialize state with database session and engine
|
|
initial_state = State(db_session=db_session, engine=engine)
|
|
|
|
# Run the graph directly
|
|
print("\nRunning the complete researcher workflow...")
|
|
result = await graph.ainvoke(initial_state, config)
|
|
|
|
# Extract the answer outline for display
|
|
if "answer_outline" in result and result["answer_outline"]:
|
|
print(f"\nGenerated answer outline with {len(result['answer_outline'].answer_outline)} sections")
|
|
|
|
# Print the outline
|
|
print("\nGenerated Answer Outline:")
|
|
for section in result["answer_outline"].answer_outline:
|
|
print(f"\nSection {section.section_id}: {section.section_title}")
|
|
print("Research Questions:")
|
|
for q in section.questions:
|
|
print(f" - {q}")
|
|
|
|
# Check if we got a final report
|
|
if "final_written_report" in result and result["final_written_report"]:
|
|
final_report = result["final_written_report"]
|
|
print("\nFinal Research Report generated successfully!")
|
|
print(f"Report length: {len(final_report)} characters")
|
|
|
|
# Display the final report
|
|
print("\n==== FINAL RESEARCH REPORT ====\n")
|
|
print(final_report)
|
|
else:
|
|
print("\nNo final report was generated.")
|
|
print(f"Available result keys: {list(result.keys())}")
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Error running researcher agent: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
async def main():
|
|
"""Main entry point for the test script."""
|
|
try:
|
|
result = await run_test()
|
|
print("\nTest completed successfully.")
|
|
return result
|
|
except Exception as e:
|
|
print(f"\nTest failed with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
if __name__ == "__main__":
|
|
# Run the async test
|
|
result = asyncio.run(main())
|
|
|
|
# Keep terminal open if run directly in VSCode
|
|
if 'VSCODE_PID' in os.environ:
|
|
input("\nPress Enter to close this window...") |