create compatible class structure

This commit is contained in:
AI Christianson 2024-12-28 14:02:52 -05:00
parent 377a670ac8
commit c704117444
1 changed files with 135 additions and 157 deletions

View File

@ -1,8 +1,11 @@
import os
import uuid
from dotenv import load_dotenv
from ra_aid.agent_utils import run_agent_with_retry
from typing import Dict, Any, Generator, List, Optional
from langchain_core.messages import AIMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from ra_aid.tools.list_directory import list_directory_tree
from ra_aid.tool_configs import get_read_only_tools
import inspect
@ -15,44 +18,14 @@ console = Console()
# Load environment variables
load_dotenv()
def get_function_info(func):
"""
Returns a well-formatted string containing the function signature and docstring,
designed to be easily readable by both humans and LLMs.
"""
# Get signature
signature = inspect.signature(func)
# Get docstring - use getdoc to clean up indentation
docstring = inspect.getdoc(func)
if docstring is None:
docstring = "No docstring provided"
# Format full signature including return type
full_signature = f"{func.__name__}{signature}"
# Build the complete string
info = f"""{full_signature}
\"\"\"
{docstring}
\"\"\" """
return info
@tool
def check_weather(location: str) -> str:
"""
Gets the weather at the given location.
"""
"""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.
"""
print()
"""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()
@ -60,78 +33,42 @@ def output_message(message: str, prompt_user_input: bool = False) -> str:
return user_input
return ""
def evaluate_response(code: str, tools: list) -> any:
class CiaynAgent:
def get_function_info(self, func):
"""
Evaluates a single function call and returns its result
Args:
code (str): The code to evaluate
tools (list): List of tool objects that have a .func property
Returns:
any: Result of the code evaluation
Returns a well-formatted string containing the function signature and docstring,
designed to be easily readable by both humans and LLMs.
"""
# Create globals dictionary from tool functions
globals_dict = {
tool.func.__name__: tool.func
for tool in tools
}
try:
# Using eval() instead of exec() since we're evaluating a single expression
result = eval(code, globals_dict)
return result
except Exception as e:
print(f"Code:\n\n{code}\n\n")
print(f"Error executing code: {str(e)}")
return f"Error executing code: {str(e)}"
def create_chat_interface():
# Initialize the chat model
chat = ChatOpenAI(
# api_key=os.getenv("OPENROUTER_API_KEY"),
api_key=os.getenv("DEEPSEEK_API_KEY"),
temperature=0.7 ,
# base_url="https://openrouter.ai/api/v1",
base_url="https://api.deepseek.com/v1",
# model="deepseek/deepseek-chat"
model="deepseek-chat"
# model="openai/gpt-4o-mini"
# model="qwen/qwen-2.5-coder-32b-instruct"
# model="qwen/qwen-2.5-72b-instruct"
)
# Chat loop
print("Welcome to the Chat Interface! (Type 'quit' to exit)")
chat_history = []
last_result = None
first_iteration = True
tools = get_read_only_tools(True, True)
tools.extend([output_message])
available_functions = []
signature = inspect.signature(func)
docstring = inspect.getdoc(func)
if docstring is None:
docstring = "No docstring provided"
full_signature = f"{func.__name__}{signature}"
info = f"""{full_signature}
\"\"\"
{docstring}
\"\"\" """
return info
def __init__(self, model, tools: list):
"""Initialize the agent with a model and list of tools."""
self.model = model
self.tools = tools
self.available_functions = []
for t in tools:
available_functions.append(get_function_info(t.func))
self.available_functions.append(self.get_function_info(t.func))
while True:
def _build_prompt(self, last_result: Optional[str] = None) -> str:
"""Build the prompt for the agent including available tools and context."""
base_prompt = ""
# Add the last result to the prompt if it's not the first iteration
if not first_iteration and last_result is not None:
if last_result is not None:
base_prompt += f"\n<last result>{last_result}</last result>"
# Construct the tool documentation and context
base_prompt += f"""
<available functions>
{"\n\n".join(available_functions)}
{"\n\n".join(self.available_functions)}
</available functions>
"""
base_prompt += """
<agent instructions>
You are a ReAct agent. You run in a loop and use ONE of the available functions per iteration.
If the current query does not require a function call, just use output_message to say what you would normally say.
@ -148,49 +85,90 @@ def create_chat_interface():
</example response>
<example response>
output_message(\"\"\"
How can I help you today?
\"\"\", True)
output_message(\"\"\"How can I help you today?\"\"\", True)
</example response>
"""
base_prompt += "\nOutput **ONLY THE CODE** and **NO MARKDOWN BACKTICKS**"
Output **ONLY THE CODE** and **NO MARKDOWN BACKTICKS**"""
return base_prompt
# Add user message to history
# Remove the previous messages if they exist
# if len(chat_history) > 1:
# chat_history.pop() # Remove the last assistant message
# chat_history.pop() # Remove the last human message
def _execute_tool(self, code: str) -> str:
"""Execute a tool call and return its result."""
globals_dict = {
tool.func.__name__: tool.func
for tool in self.tools
}
try:
result = eval(code.strip(), globals_dict)
return result
except Exception as e:
error_msg = f"Error executing code: {str(e)}"
console.print(f"[red]Error:[/red] {error_msg}")
return error_msg
def _create_agent_chunk(self, content: str) -> Dict[str, Any]:
"""Create an agent chunk in the format expected by print_agent_output."""
return {
"agent": {
"messages": [AIMessage(content=content)]
}
}
def _create_error_chunk(self, content: str) -> Dict[str, Any]:
"""Create an error chunk in the format expected by print_agent_output."""
return {
"tools": {
"messages": [{"status": "error", "content": content}]
}
}
def stream(self, messages_dict: Dict[str, List[Any]], config: Dict[str, Any] = None) -> Generator[Dict[str, Any], None, None]:
"""Stream agent responses in a format compatible with print_agent_output."""
initial_messages = messages_dict.get("messages", [])
chat_history = []
last_result = None
first_iteration = True
while True:
base_prompt = self._build_prompt(None if first_iteration else last_result)
chat_history.append(HumanMessage(content=base_prompt))
try:
# Get response from model
# print("PRECHAT")
response = chat.invoke(chat_history)
# print("POSTCHAT")
full_history = initial_messages + chat_history
response = self.model.invoke(full_history)
# # Print the code response
# print("\nAssistant generated code:")
# print(response.content)
# Evaluate the code
# print("\nExecuting code:")
# print("PREEVAL")
last_result = evaluate_response(response.content.strip(), tools)
# print("POSTEVAL")
# if last_result is not None:
# print(f"Result: {last_result}")
# Add assistant response to history
last_result = self._execute_tool(response.content)
chat_history.append(response)
# Set first_iteration to False after the first loop
first_iteration = False
# print("LOOP")
yield {}
except Exception as e:
print(f"\nError: {str(e)}")
error_msg = f"Error: {str(e)}"
yield self._create_error_chunk(error_msg)
break
if __name__ == "__main__":
create_chat_interface()
# 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)")