58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import os
|
|
import uuid
|
|
from dotenv import load_dotenv
|
|
from ra_aid.agent_utils import run_agent_with_retry
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.tools import tool
|
|
from ra_aid.tools.list_directory import list_directory_tree
|
|
from ra_aid.tool_configs import get_read_only_tools
|
|
from rich.panel import Panel
|
|
from rich.markdown import Markdown
|
|
from rich.console import Console
|
|
from ra_aid.agents.ciayn_agent import CiaynAgent
|
|
|
|
console = Console()
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
@tool
|
|
def check_weather(location: str) -> str:
|
|
"""Gets the weather at the given location."""
|
|
return f"The weather in {location} is sunny!"
|
|
|
|
@tool
|
|
def output_message(message: str, prompt_user_input: bool = False) -> str:
|
|
"""Outputs a message to the user, optionally prompting for input."""
|
|
console.print(Panel(Markdown(message.strip())))
|
|
if prompt_user_input:
|
|
user_input = input("\n> ").strip()
|
|
print()
|
|
return user_input
|
|
return ""
|
|
|
|
if __name__ == "__main__":
|
|
# Initialize the chat model
|
|
chat = ChatOpenAI(
|
|
api_key=os.getenv("OPENROUTER_API_KEY"),
|
|
temperature=0.7,
|
|
base_url="https://openrouter.ai/api/v1",
|
|
model="qwen/qwen-2.5-coder-32b-instruct"
|
|
)
|
|
|
|
# Get tools
|
|
tools = get_read_only_tools(True, True)
|
|
tools.append(output_message)
|
|
|
|
# Initialize agent
|
|
agent = CiaynAgent(chat, tools)
|
|
|
|
# Test chat prompt
|
|
test_prompt = "Find the tests in this codebase."
|
|
|
|
# Run the agent using run_agent_with_retry
|
|
result = run_agent_with_retry(agent, test_prompt, {"configurable": {"thread_id": str(uuid.uuid4())}})
|
|
|
|
# Initial greeting
|
|
print("Welcome to the Chat Interface! (Type 'quit' to exit)")
|