Spaces:
Running
Running
File size: 2,930 Bytes
016be3b |
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 |
import os
import sys
from src.research_agent import run_agent
from langchain_core.messages import AIMessage, ToolMessage
def main():
"""
Run the research agent from the command line
"""
# Check if domain is provided
if len(sys.argv) >= 2:
domain = sys.argv[1]
else:
domain = "General Research"
print(f"No domain specified, using default: {domain}")
print(f"\n=== Research Assistant ({domain}) ===\n")
print("Ask your question or type 'exit' to quit.\n")
# Main interaction loop
messages = None
while True:
# Get user input
user_input = input("> ")
# Check for exit command
if user_input.lower() in ["exit", "quit", "q"]:
print("\nGoodbye!")
break
try:
# Run the agent
print("\nResearching...\n")
messages = run_agent(user_input, domain, messages)
# Print AI response
for message in messages[-3:]: # Print only the last few messages
if hasattr(message, "type") and message.type == "ai":
print(f"\nAssistant: {message.content}\n")
elif hasattr(message, "type") and message.type == "tool" and hasattr(message, "name"):
print(f"\n[Tool: {message.name}]\n{message.content[:200]}...\n"
if len(message.content) > 200 else f"\n[Tool: {message.name}]\n{message.content}\n")
elif hasattr(message, "content"):
if isinstance(message, AIMessage):
print(f"\nAssistant: {message.content}\n")
elif isinstance(message, ToolMessage) and hasattr(message, "name"):
print(f"\n[Tool: {message.name}]\n{message.content[:200]}...\n"
if len(message.content) > 200 else f"\n[Tool: {message.name}]\n{message.content}\n")
except Exception as e:
print(f"\nError: {str(e)}\n")
print("Let's try again with a different question.")
if __name__ == "__main__":
# Check for environment variables
if not os.environ.get("OPENAI_API_KEY"):
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
pass
if not os.environ.get("OPENAI_API_KEY"):
api_key = input("Enter your OpenAI API key: ")
os.environ["OPENAI_API_KEY"] = api_key
if not os.environ.get("TAVILY_API_KEY"):
if os.path.exists(".env"):
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
pass
if not os.environ.get("TAVILY_API_KEY"):
api_key = input("Enter your Tavily API key: ")
os.environ["TAVILY_API_KEY"] = api_key
main() |