109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
import asyncio
|
|
import logging
|
|
import lmstudio as lms
|
|
from typing import Any, Callable
|
|
from .prompt import CAVEMAN_PROMPT
|
|
from .coding_prompt import CODING_AGENT_SYSTEM_PROMPT
|
|
|
|
logger: logging.Logger = logging.getLogger("agent-coding")
|
|
|
|
|
|
class _ActResponseCapture:
|
|
"""Captures the AI response from LMStudio act() callback."""
|
|
|
|
def __init__(self) -> None:
|
|
self.responses: list[str] = []
|
|
|
|
def __call__(self, message: Any) -> None:
|
|
content: str = ""
|
|
if hasattr(message, 'content'):
|
|
content = message.content
|
|
elif hasattr(message, 'text'):
|
|
content = message.text
|
|
elif hasattr(message, 'response'):
|
|
content = message.response
|
|
elif hasattr(message, 'message'):
|
|
content = message.message
|
|
else:
|
|
return
|
|
|
|
if isinstance(content, list):
|
|
parts: list[str] = []
|
|
for item in content:
|
|
if isinstance(item, dict):
|
|
text: str = item.get('text', '')
|
|
if isinstance(text, list):
|
|
parts.extend([str(t) for t in text])
|
|
else:
|
|
parts.append(str(text))
|
|
elif isinstance(item, str):
|
|
parts.append(item)
|
|
elif hasattr(item, 'text'):
|
|
parts.append(str(item.text))
|
|
elif hasattr(item, 'content'):
|
|
parts.append(str(item.content))
|
|
content = ''.join(parts)
|
|
elif not isinstance(content, str):
|
|
content = str(content)
|
|
|
|
if content.strip():
|
|
self.responses.append(content.strip())
|
|
|
|
@property
|
|
def full_response(self) -> str:
|
|
return '\n'.join(self.responses) if self.responses else "No response captured."
|
|
|
|
|
|
class CodingAgent:
|
|
"""AI agent that interacts with LMStudio models and tools."""
|
|
|
|
def __init__(self, model_name: str) -> None:
|
|
self.model_name: str = model_name
|
|
self.model: Any | None = None
|
|
self.system_prompt: str = CODING_AGENT_SYSTEM_PROMPT
|
|
|
|
async def initialize(self) -> None:
|
|
"""Initialize the LM Studio model."""
|
|
logger.info(f"Initializing CodingAgent with model: {self.model_name}")
|
|
self.model = lms.llm(self.model_name)
|
|
|
|
async def run(self, user_input: str) -> str:
|
|
"""Run a single interaction with the agent."""
|
|
if self.model is None:
|
|
await self.initialize()
|
|
assert self.model is not None
|
|
|
|
messages: list[dict[str, str]] = [
|
|
{"role": "system", "content": self.system_prompt},
|
|
{"role": "user", "content": user_input},
|
|
]
|
|
|
|
try:
|
|
logger.info(f"Running CodingAgent interactively (input length: {len(user_input)})")
|
|
response = await self.model.respond(user_input, messages=messages)
|
|
logger.info(f"CodingAgent responded successfully (response length: {len(response)})")
|
|
return response
|
|
except Exception as e:
|
|
logger.error(f"CodingAgent execution error: {e}")
|
|
return f"Error in agent execution: {str(e)}"
|
|
|
|
async def run_with_tools(self, user_input: str, tools: list[Callable[..., Any]]) -> str:
|
|
"""Run the agent with tool calling capability."""
|
|
if self.model is None:
|
|
await self.initialize()
|
|
assert self.model is not None
|
|
|
|
try:
|
|
capture = _ActResponseCapture()
|
|
logger.info(f"Calling LMStudio act() on CodingAgent with {len(tools)} tools...")
|
|
result: lms.ActResult = self.model.act(user_input, tools=tools, on_message=capture)
|
|
logger.info(f"act() on CodingAgent returned: {result}")
|
|
response: str = capture.full_response
|
|
if not response or response == "No response captured.":
|
|
logger.warning(f"Act completed with {result.rounds} rounds but no response was captured.")
|
|
return f"Act completed with {result.rounds} rounds but no response captured."
|
|
return response
|
|
except Exception as e:
|
|
logger.error(f"CodingAgent tool execution error: {e}")
|
|
return f"Error in agent tool execution: {str(e)}"
|