Untitled

import os
import json
import openai

class DummyAuditAgent:
    def __init__(self, agent_id="simple_audit_agent_v1", api_key=None):
        """
        Initializes the AuditAgent with an optional agent ID and OpenAI API key.
        """
        self.agent_id = agent_id
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        openai.api_key = self.api_key

    def analyze_contract(self, contract_code: str) -> dict:
        """
        Sends the contract code to OpenAI and attempts to parse a JSON response of findings.

        Returns:
            A dictionary that follows the format:
            {
                "agent_id": "simple_audit_agent_v1",
                "findings": [
                    {
                        "finding_id": "...",
                        "severity": "High/Medium/Low",
                        "description": "...",
                        "recommendation": "...",
                        "code_reference": "...",
                    },
                    ...
                ],
                "metadata": { ... }
            }
        """
        # Prompt instructing the AI to return valid JSON with a specific schema:
        prompt = f"""
You are a security analysis AI. Analyze the following smart contract code and return any findings in a JSON structure.
The JSON must include:
- agent_id (string)
- findings (array of objects), each with:
  - finding_id (string)
  - severity (string; e.g. "High", "Medium", or "Low")
  - description (string)
  - recommendation (string)
  - code_reference (string)
- metadata (object) for any extra info

Contract code:
{contract_code}

Ensure the response is valid JSON with no additional commentary.
        """

        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": "You are a security analysis expert."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0
            )

            # Extract the text from the response
            raw_text = response.choices[0].message.content.strip()
            
            # Attempt to parse JSON
            parsed_response = json.loads(raw_text)

        except json.JSONDecodeError:
            # If JSON parsing fails, wrap the raw output in a fallback structure
            parsed_response = {
                "agent_id": self.agent_id,
                "findings": [],
                "metadata": {
                    "error": "Unable to parse JSON from AI response",
                    "raw_response": raw_text
                }
            }
        except Exception as e:
            # Catch any other errors (e.g., network, OpenAI issues)
            parsed_response = {
                "agent_id": self.agent_id,
                "findings": [],
                "metadata": {
                    "error": str(e)
                }
            }

        # Optionally, we can also ensure that the agent_id is set correctly
        if isinstance(parsed_response, dict):
            parsed_response.setdefault("agent_id", self.agent_id)
        
        return parsed_response


if __name__ == "__main__":
    # Example usage:
    agent = DummyAuditAgent(api_key="YOUR_OPENAI_API_KEY")
    sample_contract = """
    pragma solidity ^0.8.0;

    contract Example {
        uint256 public value;

        function setValue(uint256 newValue) public {
            // No input validation
            value = newValue;
        }
    }
    """

    # Analyze the contract
    results = agent.analyze_contract(sample_contract)

    # Print the parsed JSON findings
    print(json.dumps(results, indent=2))