152 lines
5.1 KiB
Python
152 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
||
"""nanocode-dspy - minimal claude code alternative using DSPy ReAct"""
|
||
|
||
import os
|
||
from modaic import PrecompiledProgram, PrecompiledConfig
|
||
import dspy
|
||
from dspy.utils.callback import BaseCallback
|
||
|
||
from utils import (
|
||
RESET, BOLD, DIM, BLUE, CYAN, GREEN, RED, MAGENTA,
|
||
separator, render_markdown,
|
||
read_file, write_file, edit_file, glob_files, grep_files, run_bash,
|
||
AVAILABLE_MODELS, select_model,
|
||
)
|
||
|
||
|
||
# --- DSPy Signature ---
|
||
|
||
class CodingAssistant(dspy.Signature):
|
||
"""You are a concise coding assistant. Help the user with their coding task by using the available tools to read, write, edit files, search the codebase, and run commands."""
|
||
|
||
task: str = dspy.InputField(desc="The user's coding task or question")
|
||
answer: str = dspy.OutputField(desc="Your response to the user after completing the task")
|
||
affected_files: list[str] = dspy.OutputField(desc="List of files that were written or modified during the task")
|
||
|
||
# ReAct agent with tools
|
||
|
||
tools = [read_file, write_file, edit_file, glob_files, grep_files, run_bash]
|
||
|
||
|
||
class ToolLoggingCallback(BaseCallback):
|
||
"""Callback that logs tool calls as they happen."""
|
||
|
||
def on_tool_start(self, call_id, instance, inputs):
|
||
"""Log when a tool starts executing."""
|
||
tool_name = instance.name if hasattr(instance, 'name') else str(instance)
|
||
# Format args nicely
|
||
args_str = ", ".join(f"{k}={repr(v)[:50]}" for k, v in inputs.items())
|
||
print(f" {MAGENTA}⏺ {tool_name}({args_str}){RESET}", flush=True)
|
||
|
||
def on_tool_end(self, call_id, outputs, exception):
|
||
"""Log when a tool finishes executing."""
|
||
if exception:
|
||
print(f" {RED}Error: {exception}{RESET}", flush=True)
|
||
|
||
def on_module_end(self, call_id, outputs, exception):
|
||
"""Log when the finish tool is called (ReAct completion)."""
|
||
# Check if this is a ReAct prediction with tool_calls
|
||
if outputs and 'tool_calls' in outputs:
|
||
for call in outputs['tool_calls']:
|
||
args_str = ", ".join(f"{k}={repr(v)[:50]}" for k, v in call.args.items())
|
||
if call.name == 'finish':
|
||
print(f" {GREEN}⏺ finish{RESET}", flush=True)
|
||
else:
|
||
print(f" {MAGENTA}⏺ {call.name}({args_str}){RESET}", flush=True)
|
||
|
||
|
||
class AgentConfig(PrecompiledConfig):
|
||
max_iters: int = 15
|
||
lm: str = "openrouter/anthropic/claude-3.5-sonnet" # Default fallback
|
||
api_base: str = "https://openrouter.ai/api/v1"
|
||
max_tokens: int = 8192
|
||
|
||
class AgentProgram(PrecompiledProgram):
|
||
config: AgentConfig
|
||
|
||
def __init__(self, config: AgentConfig, **kwargs):
|
||
self.config = config
|
||
super().__init__(config, **kwargs)
|
||
|
||
# Configure logging callback globally
|
||
dspy.settings.configure(callbacks=[ToolLoggingCallback()])
|
||
|
||
agent = dspy.ReAct(CodingAssistant, tools=tools, max_iters=self.config.max_iters)
|
||
lm = dspy.LM(self.config.lm, api_base=self.config.api_base, max_tokens=self.config.max_tokens)
|
||
agent.set_lm(lm)
|
||
self.agent = agent
|
||
|
||
def forward(self, task: str) -> str:
|
||
assert task, "Task cannot be empty"
|
||
return self.agent(task=task)
|
||
|
||
# --- Main ---
|
||
|
||
|
||
def main():
|
||
"""Create AgentConfig with selected model."""
|
||
model = os.getenv("MODEL")
|
||
if model is None:
|
||
model = select_model()
|
||
|
||
# Add openrouter/ prefix if not already present
|
||
if not model.startswith("openrouter/"):
|
||
model = f"openrouter/{model}"
|
||
|
||
config = AgentConfig()
|
||
config.lm = model
|
||
|
||
agent = AgentProgram(config)
|
||
print(f"{BOLD}nanocode-dspy{RESET} | {DIM}{agent.config.lm} | {os.getcwd()}{RESET}\n")
|
||
|
||
# Conversation history for context
|
||
history = []
|
||
|
||
while True:
|
||
try:
|
||
print(separator())
|
||
user_input = input(f"{BOLD}{BLUE}❯{RESET} ").strip()
|
||
print(separator())
|
||
|
||
if not user_input:
|
||
continue
|
||
if user_input in ("/q", "exit"):
|
||
break
|
||
if user_input == "/c":
|
||
history = []
|
||
print(f"{GREEN}⏺ Cleared conversation{RESET}")
|
||
continue
|
||
|
||
# Build context from history
|
||
context = f"Working directory: {os.getcwd()}\n"
|
||
if history:
|
||
context += "\nPrevious conversation:\n"
|
||
for h in history[-5:]: # Keep last 5 exchanges
|
||
context += f"User: {h['user']}\nAssistant: {h['assistant']}\n\n"
|
||
|
||
task = f"{context}\nCurrent task: {user_input}"
|
||
|
||
print(f"\n{CYAN}⏺{RESET} Thinking...", flush=True)
|
||
|
||
# Run the ReAct agent
|
||
result = agent(task=task)
|
||
|
||
# Display the answer
|
||
print(f"\n{CYAN}⏺{RESET} {render_markdown(result.answer)}")
|
||
|
||
# Save to history
|
||
history.append({"user": user_input, "assistant": result.answer})
|
||
|
||
print()
|
||
|
||
except (KeyboardInterrupt, EOFError):
|
||
break
|
||
except Exception as err:
|
||
import traceback
|
||
traceback.print_exc()
|
||
print(f"{RED}⏺ Error: {err}{RESET}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|