Final fixes + test + cleanup
This commit is contained in:
parent
b096096f93
commit
65ab978e68
2
.gitignore
vendored
2
.gitignore
vendored
@ -165,3 +165,5 @@ config.json
|
||||
*.jsonl
|
||||
example_session.jsonl
|
||||
|
||||
*.pyc
|
||||
__pycache__/
|
||||
|
||||
@ -1,11 +1,10 @@
|
||||
"""
|
||||
PicoGent - Minimal AI Coding Agent
|
||||
PicoGent - Minimal AI coding agent
|
||||
"""
|
||||
|
||||
from .agent import Agent, run_agent
|
||||
from .config import Config
|
||||
from .session import Session
|
||||
from .context import Context
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__all__ = ["Agent", "run_agent", "Config", "Session", "Context"]
|
||||
|
||||
from .agent import Agent
|
||||
from .config import Config
|
||||
|
||||
__all__ = ["Agent", "Config"]
|
||||
@ -1,69 +1,72 @@
|
||||
"""
|
||||
Context management for PicoGent
|
||||
Context builder for system prompt, memory, and skills
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, Any, List, Optional
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
class Context:
|
||||
"""Context manager for agent workspace and environment"""
|
||||
class ContextBuilder:
|
||||
"""Build context from system prompt, memory, and skills"""
|
||||
|
||||
def __init__(self, workspace: str = "."):
|
||||
self.workspace = os.path.abspath(workspace)
|
||||
self.context_data: Dict[str, Any] = {}
|
||||
def __init__(self, workspace_path: str = "."):
|
||||
self.workspace_path = workspace_path
|
||||
|
||||
def set_workspace(self, workspace: str):
|
||||
"""Set the workspace directory"""
|
||||
self.workspace = os.path.abspath(workspace)
|
||||
if not os.path.exists(self.workspace):
|
||||
os.makedirs(self.workspace, exist_ok=True)
|
||||
|
||||
def get_workspace(self) -> str:
|
||||
"""Get the current workspace directory"""
|
||||
return self.workspace
|
||||
|
||||
def resolve_path(self, path: str) -> str:
|
||||
"""Resolve a path relative to workspace"""
|
||||
if os.path.isabs(path):
|
||||
return path
|
||||
return os.path.join(self.workspace, path)
|
||||
|
||||
def set_context(self, key: str, value: Any):
|
||||
"""Set context data"""
|
||||
self.context_data[key] = value
|
||||
|
||||
def get_context(self, key: str, default: Any = None) -> Any:
|
||||
"""Get context data"""
|
||||
return self.context_data.get(key, default)
|
||||
|
||||
def clear_context(self):
|
||||
"""Clear all context data"""
|
||||
self.context_data.clear()
|
||||
|
||||
def get_system_prompt(self) -> str:
|
||||
"""Generate system prompt with context"""
|
||||
prompt_parts = [
|
||||
"You are PicoGent, a helpful coding assistant.",
|
||||
f"Your workspace is: {self.workspace}",
|
||||
"",
|
||||
"Available tools:",
|
||||
"- read: Read file contents with optional offset/limit",
|
||||
"- write: Write content to files (creates directories as needed)",
|
||||
"- edit: Replace text in files (old_string -> new_string)",
|
||||
"- bash: Execute shell commands (30s timeout)",
|
||||
"",
|
||||
"Guidelines:",
|
||||
"1. Always use relative paths when possible",
|
||||
"2. Be precise with file operations",
|
||||
"3. Ask for confirmation before destructive operations",
|
||||
"4. Provide helpful error messages",
|
||||
"5. Keep responses concise and actionable"
|
||||
]
|
||||
def build_system_prompt(self, base_prompt: str) -> str:
|
||||
"""Build system prompt with skills and context"""
|
||||
system_parts = [base_prompt]
|
||||
|
||||
# Add custom context if available
|
||||
custom_context = self.get_context("system_additions")
|
||||
if custom_context:
|
||||
prompt_parts.extend(["", "Additional context:", custom_context])
|
||||
# Add skills from skills directory
|
||||
skills_dir = os.path.join(self.workspace_path, "skills")
|
||||
if os.path.exists(skills_dir):
|
||||
skills = self._load_skills(skills_dir)
|
||||
if skills:
|
||||
system_parts.append("\n## Available Skills\n")
|
||||
system_parts.extend(skills)
|
||||
|
||||
return "\n".join(prompt_parts)
|
||||
return "\n".join(system_parts)
|
||||
|
||||
def _load_skills(self, skills_dir: str) -> List[str]:
|
||||
"""Load skills from skills directory"""
|
||||
skills = []
|
||||
|
||||
if not os.path.exists(skills_dir):
|
||||
return skills
|
||||
|
||||
for filename in sorted(os.listdir(skills_dir)):
|
||||
if filename.endswith(('.md', '.txt')):
|
||||
skill_path = os.path.join(skills_dir, filename)
|
||||
try:
|
||||
with open(skill_path, 'r', encoding='utf-8') as f:
|
||||
skill_content = f.read().strip()
|
||||
if skill_content:
|
||||
skills.append(f"### {filename}\n{skill_content}")
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not load skill {filename}: {e}")
|
||||
|
||||
return skills
|
||||
|
||||
def get_working_directory_info(self) -> str:
|
||||
"""Get information about current working directory"""
|
||||
try:
|
||||
files = []
|
||||
dirs = []
|
||||
|
||||
for item in os.listdir(self.workspace_path):
|
||||
item_path = os.path.join(self.workspace_path, item)
|
||||
if os.path.isdir(item_path):
|
||||
dirs.append(item + "/")
|
||||
else:
|
||||
files.append(item)
|
||||
|
||||
info_parts = [f"Working directory: {os.path.abspath(self.workspace_path)}"]
|
||||
if dirs or files:
|
||||
all_items = sorted(dirs) + sorted(files)
|
||||
info_parts.append("Contents: " + ", ".join(all_items[:20])) # Limit to first 20 items
|
||||
if len(all_items) > 20:
|
||||
info_parts.append(f"... and {len(all_items) - 20} more items")
|
||||
|
||||
return "\n".join(info_parts)
|
||||
|
||||
except Exception as e:
|
||||
return f"Working directory: {os.path.abspath(self.workspace_path)} (could not list contents: {e})"
|
||||
@ -1,118 +1,87 @@
|
||||
"""
|
||||
Anthropic provider using httpx
|
||||
Anthropic Claude provider using httpx
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import json
|
||||
import uuid
|
||||
from typing import List, Dict, Any, Optional
|
||||
from typing import Dict, List, Any, Optional
|
||||
from .base import BaseProvider
|
||||
|
||||
|
||||
class AnthropicProvider(BaseProvider):
|
||||
"""Anthropic provider using httpx directly"""
|
||||
"""Anthropic Claude provider using direct API calls"""
|
||||
|
||||
API_BASE = "https://api.anthropic.com/v1"
|
||||
API_VERSION = "2023-06-01"
|
||||
def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514"):
|
||||
super().__init__(api_key, model)
|
||||
self.base_url = "https://api.anthropic.com/v1"
|
||||
self.anthropic_version = "2023-06-01"
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.client = httpx.AsyncClient()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
await self.client.aclose()
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
async def generate_response(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
system_prompt: str,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system_prompt: Optional[str] = None
|
||||
max_tokens: int = 8192
|
||||
) -> Dict[str, Any]:
|
||||
"""Get chat completion from Anthropic API"""
|
||||
"""Generate response using Anthropic Messages API"""
|
||||
|
||||
headers = {
|
||||
"x-api-key": self.config.api_key,
|
||||
"anthropic-version": self.API_VERSION,
|
||||
"content-type": "application/json"
|
||||
"Content-Type": "application/json",
|
||||
"x-api-key": self.api_key,
|
||||
"anthropic-version": self.anthropic_version
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": self.config.model,
|
||||
"max_tokens": self.config.max_tokens,
|
||||
"model": self.model,
|
||||
"max_tokens": max_tokens,
|
||||
"system": system_prompt,
|
||||
"messages": messages
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
|
||||
# Add tools if provided
|
||||
if tools:
|
||||
payload["tools"] = tools
|
||||
|
||||
try:
|
||||
response = await self.client.post(
|
||||
f"{self.API_BASE}/messages",
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/messages",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=60.0
|
||||
timeout=120.0
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
error_text = response.text
|
||||
return {
|
||||
"error": f"API Error {response.status_code}: {error_text}",
|
||||
"content": f"Error: Anthropic API returned {response.status_code}"
|
||||
}
|
||||
raise Exception(f"Anthropic API error {response.status_code}: {error_text}")
|
||||
|
||||
result = response.json()
|
||||
|
||||
# Extract content and tool calls
|
||||
content_blocks = result.get("content", [])
|
||||
text_content = ""
|
||||
# Parse the response
|
||||
content = result.get("content", [])
|
||||
|
||||
# Extract text content and tool calls
|
||||
text_parts = []
|
||||
tool_calls = []
|
||||
|
||||
for block in content_blocks:
|
||||
for block in content:
|
||||
if block.get("type") == "text":
|
||||
text_content += block.get("text", "")
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif block.get("type") == "tool_use":
|
||||
tool_calls.append({
|
||||
"id": block.get("id"),
|
||||
"name": block.get("name"),
|
||||
"arguments": block.get("input", {})
|
||||
"input": block.get("input", {})
|
||||
})
|
||||
|
||||
return {
|
||||
"content": text_content,
|
||||
"tool_calls": tool_calls,
|
||||
"raw_response": result
|
||||
response_data = {
|
||||
"content": "\n".join(text_parts) if text_parts else "",
|
||||
"usage": result.get("usage", {}),
|
||||
"model": result.get("model", self.model)
|
||||
}
|
||||
|
||||
except httpx.TimeoutException:
|
||||
return {
|
||||
"error": "Request timed out",
|
||||
"content": "Error: Request to Anthropic API timed out"
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"error": str(e),
|
||||
"content": f"Error: {str(e)}"
|
||||
}
|
||||
|
||||
def parse_tool_calls(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""Parse tool calls from Anthropic response"""
|
||||
return response.get("tool_calls", [])
|
||||
|
||||
def format_tool_result(self, tool_call_id: str, result: str) -> Dict[str, Any]:
|
||||
"""Format tool result for Anthropic API"""
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tool_call_id,
|
||||
"content": result
|
||||
}
|
||||
]
|
||||
}
|
||||
if tool_calls:
|
||||
response_data["tool_calls"] = tool_calls
|
||||
# Also include raw content for session storage
|
||||
response_data["raw_content"] = content
|
||||
|
||||
return response_data
|
||||
65
test_basic.py
Normal file
65
test_basic.py
Normal file
@ -0,0 +1,65 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Basic functionality test for PicoGent
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
import os
|
||||
from picogent import Config, Session, run_agent
|
||||
|
||||
|
||||
async def test_basic_functionality():
|
||||
"""Test basic agent functionality"""
|
||||
print("🧪 Testing PicoGent basic functionality...")
|
||||
|
||||
# Create a test config
|
||||
config = Config(
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4-20250514",
|
||||
api_key="test-key", # Fake key for testing
|
||||
max_tokens=1000,
|
||||
max_iterations=3,
|
||||
workspace="."
|
||||
)
|
||||
|
||||
print("✅ Config created successfully")
|
||||
|
||||
# Test session creation
|
||||
session = Session()
|
||||
session.add_message("user", "Hello, world!")
|
||||
session.add_message("assistant", "Hello! How can I help you?")
|
||||
|
||||
messages = session.get_messages()
|
||||
assert len(messages) == 2
|
||||
assert messages[0]["role"] == "user"
|
||||
assert messages[1]["role"] == "assistant"
|
||||
|
||||
print("✅ Session management working")
|
||||
|
||||
# Test tools import and registration
|
||||
from picogent.tools import registry, ReadTool, WriteTool, EditTool, BashTool
|
||||
|
||||
# Should have some tools registered
|
||||
tools = registry.list_tools()
|
||||
print(f"✅ {len(tools)} tools registered: {[t.name for t in tools]}")
|
||||
|
||||
# Test tool definitions
|
||||
definitions = registry.get_tool_definitions()
|
||||
assert len(definitions) > 0
|
||||
assert all("name" in tool for tool in definitions)
|
||||
|
||||
print("✅ Tool definitions generated successfully")
|
||||
|
||||
# Test providers
|
||||
from picogent.providers import AnthropicProvider
|
||||
provider = AnthropicProvider(config)
|
||||
|
||||
print("✅ Anthropic provider created successfully")
|
||||
|
||||
print("✅ All basic tests passed!")
|
||||
print("\n📝 Note: To test actual API functionality, set a valid ANTHROPIC_API_KEY environment variable")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_basic_functionality())
|
||||
Loading…
Reference in New Issue
Block a user