285 lines
11 KiB
Python
285 lines
11 KiB
Python
|
from typing import List, Any, Dict, Literal, Optional
|
||
|
|
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
|
||
|
|
from langgraph.graph import END
|
||
|
|
from agents.base import BaseAgent, AgentState
|
||
|
|
from tools.file.write_to_file import write_to_file
|
||
|
|
from tools.agent.assign_agent_to_task import assign_agent_to_task
|
||
|
|
from tools.file.delete_file import delete_file
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import logging
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
# Configure logging
|
||
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
class TechSageAgent(BaseAgent):
|
||
|
|
def __init__(self, max_iterations: int = 1):
|
||
|
|
system_prompt = """You are techsage, a specialized development agent.
|
||
|
|
|
||
|
|
Your primary task is to generate well-structured, production-ready code based on user requirements.
|
||
|
|
|
||
|
|
Response Format:
|
||
|
|
Always structure your response with:
|
||
|
|
1. Code implementation in clear, distinct sections with proper file markers
|
||
|
|
2. Each file should be marked with ```language filename.ext
|
||
|
|
3. Include setup instructions and usage examples
|
||
|
|
4. Provide clear API documentation if applicable
|
||
|
|
|
||
|
|
Follow these guidelines:
|
||
|
|
1. Use modern best practices
|
||
|
|
2. Include error handling
|
||
|
|
3. Add proper comments and documentation
|
||
|
|
4. Optimize for readability and maintainability
|
||
|
|
5. Consider scalability and performance
|
||
|
|
|
||
|
|
Code Structure:
|
||
|
|
- Organize code logically
|
||
|
|
- Use consistent formatting
|
||
|
|
- Include necessary imports
|
||
|
|
- Add type hints where applicable
|
||
|
|
- Implement error handling
|
||
|
|
- Add proper validation
|
||
|
|
|
||
|
|
Documentation Guidelines:
|
||
|
|
- Clear setup instructions
|
||
|
|
- Usage examples
|
||
|
|
- API documentation
|
||
|
|
- Configuration details
|
||
|
|
- Dependencies list
|
||
|
|
- Error handling guide"""
|
||
|
|
|
||
|
|
super().__init__(
|
||
|
|
name="techsage",
|
||
|
|
system_prompt=system_prompt,
|
||
|
|
tools=[write_to_file, assign_agent_to_task],
|
||
|
|
max_iterations=max_iterations
|
||
|
|
)
|
||
|
|
|
||
|
|
def _analyze_task(self, task: str) -> Dict[str, Any]:
|
||
|
|
"""Analyze the task to determine type and requirements with improved error handling."""
|
||
|
|
analysis_prompt = f"""Analyze this development task and provide structured output:
|
||
|
|
Task: {task}
|
||
|
|
|
||
|
|
Format response as a valid JSON with these keys:
|
||
|
|
- task_type: web/script/config/documentation
|
||
|
|
- language: programming language name
|
||
|
|
- files_required: list of filenames
|
||
|
|
- technologies: relevant technologies
|
||
|
|
- implementation_approach: brief strategy
|
||
|
|
- primary_features: key features"""
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = self.llm.invoke(analysis_prompt)
|
||
|
|
# Enhanced parsing to handle various JSON formats
|
||
|
|
content = response.content.strip()
|
||
|
|
|
||
|
|
# Remove code block markers if present
|
||
|
|
content = re.sub(r'^```(json)?|```$', '', content, flags=re.MULTILINE).strip()
|
||
|
|
|
||
|
|
# Attempt to parse JSON with fallback
|
||
|
|
try:
|
||
|
|
analysis = json.loads(content)
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
# Attempt to fix common JSON formatting issues
|
||
|
|
content = re.sub(r'(?<=\w)\'', '"', content) # Replace single quotes with double quotes
|
||
|
|
analysis = json.loads(content)
|
||
|
|
|
||
|
|
return analysis
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Task analysis error: {e}")
|
||
|
|
return {
|
||
|
|
"task_type": "script",
|
||
|
|
"language": "python",
|
||
|
|
"files_required": ["main.py"],
|
||
|
|
"technologies": ["python"],
|
||
|
|
"implementation_approach": "Basic implementation",
|
||
|
|
"primary_features": ["core functionality"]
|
||
|
|
}
|
||
|
|
|
||
|
|
def _get_implementation_prompt(self, task: str, analysis: Dict[str, Any]) -> str:
|
||
|
|
"""Generate a more comprehensive implementation prompt."""
|
||
|
|
return f"""Comprehensive Code Generation Task
|
||
|
|
|
||
|
|
Detailed Requirements:
|
||
|
|
- Primary Task: {task}
|
||
|
|
- Language: {analysis['language']}
|
||
|
|
- Project Type: {analysis['task_type']}
|
||
|
|
- Key Features: {', '.join(analysis['primary_features'])}
|
||
|
|
|
||
|
|
Comprehensive Implementation Guidelines:
|
||
|
|
1. Create full implementation for each required file
|
||
|
|
2. Use modern {analysis['language']} best practices
|
||
|
|
3. Include robust error handling
|
||
|
|
4. Implement input validation
|
||
|
|
5. Add comprehensive type hints
|
||
|
|
6. Write clear, explanatory comments
|
||
|
|
|
||
|
|
Structural Requirements:
|
||
|
|
- Each file must be marked with: ```{analysis['language']} filename.ext
|
||
|
|
- Include complete implementation
|
||
|
|
- Add section headers for:
|
||
|
|
a. Setup Instructions
|
||
|
|
b. Usage Examples
|
||
|
|
c. API Documentation (if applicable)
|
||
|
|
d. Configuration Guide
|
||
|
|
e. Error Handling Guide
|
||
|
|
|
||
|
|
Provide a production-ready solution that emphasizes:
|
||
|
|
- Code quality
|
||
|
|
- Maintainability
|
||
|
|
- Scalability
|
||
|
|
- Performance considerations"""
|
||
|
|
|
||
|
|
def _extract_code_blocks(self, content: str) -> Dict[str, Dict[str, Any]]:
|
||
|
|
"""Enhanced code block extraction with robust parsing."""
|
||
|
|
code_blocks = {}
|
||
|
|
code_block_pattern = re.compile(r'```(\w+)?\s*(\S+)\n(.*?)```', re.DOTALL)
|
||
|
|
|
||
|
|
for match in code_block_pattern.finditer(content):
|
||
|
|
language = match.group(1) or 'text'
|
||
|
|
filename = match.group(2)
|
||
|
|
code = match.group(3).strip()
|
||
|
|
|
||
|
|
code_blocks[filename] = {
|
||
|
|
"language": language,
|
||
|
|
"content": code,
|
||
|
|
"filename": filename
|
||
|
|
}
|
||
|
|
|
||
|
|
return code_blocks
|
||
|
|
|
||
|
|
def _process_step(self, state: AgentState) -> AgentState:
|
||
|
|
"""Process a development task with enhanced error handling and logging."""
|
||
|
|
logger.info(f"{self.name} is processing development task...")
|
||
|
|
messages: List[BaseMessage] = state.get('messages', [])
|
||
|
|
|
||
|
|
try:
|
||
|
|
task = messages[-1].content if messages and hasattr(messages[-1], 'content') else "No task provided"
|
||
|
|
|
||
|
|
# Task analysis
|
||
|
|
analysis = self._analyze_task(task)
|
||
|
|
|
||
|
|
# Implementation generation
|
||
|
|
implementation_prompt = self._get_implementation_prompt(task, analysis)
|
||
|
|
implementation = self.llm.invoke(implementation_prompt)
|
||
|
|
|
||
|
|
# Process code blocks
|
||
|
|
code_blocks = self._extract_code_blocks(implementation.content)
|
||
|
|
|
||
|
|
# Write files
|
||
|
|
files_created = []
|
||
|
|
for filename, file_data in code_blocks.items():
|
||
|
|
try:
|
||
|
|
write_to_file.invoke({
|
||
|
|
"filepath": filename,
|
||
|
|
"content": file_data["content"]
|
||
|
|
})
|
||
|
|
files_created.append(filename)
|
||
|
|
except Exception as write_error:
|
||
|
|
logger.error(f"File write error for {filename}: {write_error}")
|
||
|
|
|
||
|
|
# Comprehensive response generation
|
||
|
|
response_data = {
|
||
|
|
"status": "success",
|
||
|
|
"query": task,
|
||
|
|
"timestamp": datetime.now().isoformat(),
|
||
|
|
"analysis": {
|
||
|
|
"task_type": analysis.get("task_type", "undefined"),
|
||
|
|
"language": analysis.get("language", "undefined"),
|
||
|
|
"technologies": analysis.get("technologies", [])
|
||
|
|
},
|
||
|
|
"implementation": {
|
||
|
|
"files": [
|
||
|
|
{
|
||
|
|
"filename": file_data["filename"],
|
||
|
|
"language": file_data["language"],
|
||
|
|
"content": file_data["content"]
|
||
|
|
}
|
||
|
|
for file_data in code_blocks.values()
|
||
|
|
],
|
||
|
|
"setup": self._extract_section(implementation.content, "Setup Instructions"),
|
||
|
|
"usage": self._extract_section(implementation.content, "Usage Examples"),
|
||
|
|
"api_docs": self._extract_section(implementation.content, "API Documentation"),
|
||
|
|
"configuration": self._extract_section(implementation.content, "Configuration Guide")
|
||
|
|
},
|
||
|
|
"files_created": files_created
|
||
|
|
}
|
||
|
|
|
||
|
|
for filename in files_created:
|
||
|
|
try:
|
||
|
|
delete_file.invoke({"filepath": filename})
|
||
|
|
except Exception as delete_error:
|
||
|
|
logger.error(f"File deletion error for {filename}: {delete_error}")
|
||
|
|
|
||
|
|
return {
|
||
|
|
"messages": messages + [AIMessage(content=json.dumps(response_data, indent=2))],
|
||
|
|
"iterations": 1
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Development task processing error: {e}")
|
||
|
|
error_response = {
|
||
|
|
"status": "error",
|
||
|
|
"query": task,
|
||
|
|
"timestamp": datetime.now().isoformat(),
|
||
|
|
"error": str(e),
|
||
|
|
"message": "Comprehensive error in code generation",
|
||
|
|
"suggestions": [
|
||
|
|
"Provide more specific and granular requirements",
|
||
|
|
"Clearly specify the programming language and framework",
|
||
|
|
"Break down complex requirements into smaller, manageable tasks",
|
||
|
|
"Verify the input task description"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
return {
|
||
|
|
"messages": messages + [AIMessage(content=json.dumps(error_response, indent=2))],
|
||
|
|
"iterations": 1
|
||
|
|
}
|
||
|
|
|
||
|
|
def _extract_section(self, content: str, section_name: str) -> str:
|
||
|
|
"""Enhanced section extraction with regex and multiple parsing strategies."""
|
||
|
|
try:
|
||
|
|
# Regex pattern to find section content
|
||
|
|
section_pattern = re.compile(
|
||
|
|
rf'{section_name}:\n(.*?)(?=\n\n|\Z)',
|
||
|
|
re.DOTALL | re.IGNORECASE
|
||
|
|
)
|
||
|
|
match = section_pattern.search(content)
|
||
|
|
|
||
|
|
if match:
|
||
|
|
return match.group(1).strip()
|
||
|
|
|
||
|
|
# Fallback parsing strategy
|
||
|
|
if section_name in content:
|
||
|
|
parts = content.split(section_name)
|
||
|
|
if len(parts) > 1:
|
||
|
|
section = parts[1].split('\n\n')[0].strip()
|
||
|
|
return section
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Section extraction error for {section_name}: {e}")
|
||
|
|
|
||
|
|
return ""
|
||
|
|
|
||
|
|
def _should_continue(self, state: AgentState) -> Literal["continue", END]:
|
||
|
|
"""Always terminate after one iteration."""
|
||
|
|
return END
|
||
|
|
|
||
|
|
def techsage(task: str) -> str:
|
||
|
|
"""Execute development task and return comprehensive results."""
|
||
|
|
try:
|
||
|
|
agent = TechSageAgent()
|
||
|
|
return agent.process(task)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Tech Sage agent execution failed: {e}")
|
||
|
|
return json.dumps({
|
||
|
|
"status": "critical_error",
|
||
|
|
"message": "Failed to execute engineering task",
|
||
|
|
"error": str(e),
|
||
|
|
"suggestions": [
|
||
|
|
"Retry the task",
|
||
|
|
"Verify input requirements",
|
||
|
|
"Contact system administrator"
|
||
|
|
]
|
||
|
|
})
|