Rob200/LegacyCodeDocumentAgent

1"""
2Core agent module for the Legacy Code Documentation Agent.
3Handles communication with the LLM.
4"""
5
6import os
7from dotenv import load_dotenv
8import litellm
9
10from prompts import SYSTEM_PROMPT, USER_PROMPT_TEMPLATE
11
12# Load environment variables
13load_dotenv()
14
15# Get the default model from .env
16DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "gpt-4o-mini")
17
18# Cost per 1K tokens (approximate, as of 2024)
19# Update these as pricing changes
20COST_PER_1K_TOKENS = {
21    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
22    "gpt-4o": {"input": 0.005, "output": 0.015},
23    "gemini/gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
24    "gemini/gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
25    "claude-sonnet-4-20250514": {"input": 0.003, "output": 0.015},
26    "claude-opus-4-20250514": {"input": 0.015, "output": 0.075},
27}
28
29
30def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
31    """
32    Calculate the estimated cost for an API call.
33
34    Args:
35        model: The model name
36        input_tokens: Number of input tokens
37        output_tokens: Number of output tokens
38
39    Returns:
40        Estimated cost in USD
41    """
42    if model not in COST_PER_1K_TOKENS:
43        return 0.0  # Unknown model, can't estimate
44
45    pricing = COST_PER_1K_TOKENS[model]
46    input_cost = (input_tokens / 1000) * pricing["input"]
47    output_cost = (output_tokens / 1000) * pricing["output"]
48
49    return input_cost + output_cost
50
51
52def generate_documentation(filename: str, language: str, code_content: str) -> dict:
53    """
54    Send code to the LLM and get documentation back.
55
56    Args:
57        filename: Name of the code file
58        language: Programming language (SQL, Python, etc.)
59        code_content: The actual code to document
60
61    Returns:
62        Dictionary with keys: 'success', 'documentation', 'error', 'model_used', 'usage'
63    """
64    result = {
65        "success": False,
66        "documentation": None,
67        "error": None,
68        "model_used": DEFAULT_MODEL,
69        "usage": {
70            "input_tokens": 0,
71            "output_tokens": 0,
72            "total_tokens": 0,
73            "estimated_cost": 0.0
74        }
75    }
76
77    # Build the user message from our template
78    user_message = USER_PROMPT_TEMPLATE.format(
79        language=language,
80        filename=filename,
81        code_content=code_content
82    )
83
84    try:
85        # Make the API call using litellm
86        response = litellm.completion(
87            model=DEFAULT_MODEL,
88            messages=[
89                {"role": "system", "content": SYSTEM_PROMPT},
90                {"role": "user", "content": user_message}
91            ]
92        )
93
94        # Extract the documentation from the response
95        result["documentation"] = response.choices[0].message.content
96        result["success"] = True
97
98        # Capture token usage
99        if hasattr(response, 'usage') and response.usage:
100            result["usage"]["input_tokens"] = response.usage.prompt_tokens
101            result["usage"]["output_tokens"] = response.usage.completion_tokens
102            result["usage"]["total_tokens"] = response.usage.total_tokens
103            result["usage"]["estimated_cost"] = calculate_cost(
104                DEFAULT_MODEL,
105                response.usage.prompt_tokens,
106                response.usage.completion_tokens
107            )
108
109    except Exception as e:
110        result["error"] = f"LLM API error: {str(e)}"
111
112    return result