116 lines
No EOL
4.2 KiB
Python
116 lines
No EOL
4.2 KiB
Python
# Example agent implementation
|
|
# This demonstrates the integration pattern with AI Skills API
|
|
|
|
import os
|
|
import asyncio
|
|
import httpx
|
|
from typing import List, Dict, Optional
|
|
|
|
API_URL = os.getenv("API_URL", "http://helm:8675")
|
|
API_KEY = os.getenv("API_KEY")
|
|
|
|
async def get_context(query: str, project: Optional[str] = None) -> Dict:
|
|
"""Fetch relevant context from skills API"""
|
|
params = {"query": query}
|
|
if project:
|
|
params["project"] = project
|
|
|
|
headers = {"X-API-Key": API_KEY} if API_KEY else {}
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get(f"{API_URL}/context/rag", params=params, headers=headers)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
async def compress_messages(messages: List[Dict]) -> Dict:
|
|
"""Compress conversation history"""
|
|
headers = {"X-API-Key": API_KEY} if API_KEY else {}
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(f"{API_URL}/compress", json={"messages": messages}, headers=headers)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
async def store_memory(project: str, key: str, content: str) -> Dict:
|
|
"""Store a memory for future reference"""
|
|
headers = {"X-API-Key": API_KEY} if API_KEY else {}
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.post(
|
|
f"{API_URL}/memory",
|
|
json={"id": key[:8], "project": project, "key": key, "content": content},
|
|
headers=headers
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
async def count_tokens(text: str) -> int:
|
|
"""Count tokens using skills API"""
|
|
headers = {"X-API-Key": API_KEY} if API_KEY else {}
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get(f"{API_URL}/tokens/count", params={"text": text}, headers=headers)
|
|
resp.raise_for_status()
|
|
return resp.json()["tokens"]
|
|
|
|
async def chat_loop():
|
|
"""Main chat loop - integrate with your LLM of choice"""
|
|
conversation = []
|
|
|
|
print("Agent ready! Type 'quit' to exit.")
|
|
|
|
while True:
|
|
user_input = input("\nYou: ")
|
|
if user_input.lower() == 'quit':
|
|
break
|
|
|
|
# 1. Get relevant context
|
|
context = await get_context(user_input, project="/home/user/projects/myapp")
|
|
context_str = format_context(context)
|
|
|
|
# 2. Build prompt with context
|
|
system_msg = f"{context_str}\n\nYou are a helpful assistant."
|
|
messages = [{"role": "system", "content": system_msg}]
|
|
messages.extend(conversation[-4:]) # Keep last few turns
|
|
messages.append({"role": "user", "content": user_input})
|
|
|
|
# 3. Call your LLM here (not included - use OpenAI, Claude, Ollama, etc.)
|
|
# response = await call_llm(messages)
|
|
# For demo, we'll just echo
|
|
response = f"Echo: {user_input}"
|
|
|
|
# 4. Update conversation
|
|
conversation.append({"role": "user", "content": user_input})
|
|
conversation.append({"role": "assistant", "content": response})
|
|
|
|
# 5. Compress if getting long
|
|
if len(conversation) > 10:
|
|
compression = await compress_messages(conversation)
|
|
conversation = compression["messages"]
|
|
print(f"\n[Compressed: saved {compression['tokens_saved']} tokens]")
|
|
|
|
print(f"\nAssistant: {response}")
|
|
|
|
def format_context(context: Dict) -> str:
|
|
"""Format RAG context for inclusion in prompt"""
|
|
parts = []
|
|
|
|
if context.get("skills"):
|
|
parts.append("## Relevant Skills\n")
|
|
for skill in context["skills"]:
|
|
parts.append(f"### {skill['name']} (relevance: {skill['relevance_score']:.2f})\n{skill['content']}\n")
|
|
|
|
if context.get("conventions"):
|
|
parts.append("## Project Conventions\n")
|
|
for conv in context["conventions"]:
|
|
parts.append(f"### {conv['name']}\n{conv['content']}\n")
|
|
|
|
if context.get("snippets"):
|
|
parts.append("## Code Snippets\n")
|
|
for snippet in context["snippets"]:
|
|
parts.append(f"### {snippet['name']} ({snippet['language']})\n```{snippet['language']}\n{snippet['content']}\n```\n")
|
|
|
|
return "\n".join(parts) if parts else "No relevant context found."
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(chat_loop()) |