zackeryyy wang
644071850a
- AI-powered shell command generator using DeepSeek V3 - Support for natural language to shell command conversion - Secure configuration management with .env files - Package structure with uv tool installation support - Chinese and English language support - Configuration validation and error handling
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
"""
|
|
AI Agent module for shell command generation
|
|
"""
|
|
|
|
from textwrap import dedent
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.models.openai import OpenAIModel
|
|
|
|
from .config import get_model, setup_environment
|
|
from .models import Answer
|
|
|
|
# System prompt for the AI agent
|
|
SYSTEM_PROMPT = dedent(
|
|
"""\
|
|
You are a professional developer specializing in shell commands.
|
|
Your task is to generate the correct shell commands based on the
|
|
user's request.
|
|
|
|
IMPORTANT: ALWAYS USE THE SAME LANGUAGE AS THE USER PROMPT IN
|
|
YOUR RESPONSE.
|
|
|
|
Process:
|
|
|
|
1. Think Aloud: Use the `think` function to explain your reasoning.
|
|
Justify why you chose a particular command, considering efficiency,
|
|
safety, and best practices.
|
|
|
|
2. Provide the Final Command: Use the `answer` function to present
|
|
the final shell command concisely.
|
|
"""
|
|
)
|
|
|
|
def create_agent() -> Agent:
|
|
"""Create and configure the AI agent"""
|
|
# Setup environment variables
|
|
setup_environment()
|
|
|
|
# Create OpenAI compatible model
|
|
model = OpenAIModel(get_model())
|
|
|
|
# Create agent
|
|
agent = Agent(
|
|
model=model,
|
|
system_prompt=SYSTEM_PROMPT,
|
|
output_type=Answer,
|
|
)
|
|
|
|
# Register tools
|
|
@agent.tool_plain
|
|
def think(s: str) -> None:
|
|
"""Communicate your thought process to the user.
|
|
|
|
Args:
|
|
s (str): A description of your reasoning or decision-making process.
|
|
"""
|
|
print(f"(AI Thinking): {s}\n")
|
|
|
|
@agent.tool_plain
|
|
def answer(success: bool, cmd: str | None, failure: str | None) -> Answer:
|
|
"""Provide the final shell command or explain why it couldn't be generated.
|
|
|
|
Args:
|
|
success (bool): Indicates whether a shell command was successfully generated.
|
|
cmd (str | None): The generated shell command if `success` is True.
|
|
It must be a single-line command. If `success` is False, this should be None.
|
|
failure (str | None): If `success` is False, provide a reason why the command
|
|
could not be generated. If `success` is True, this should be None.
|
|
|
|
Returns:
|
|
Answer: A structured response that will be processed for the user.
|
|
"""
|
|
return Answer(success, cmd, failure)
|
|
|
|
return agent
|