Untitled

import os
import json
import openai
from flask import Flask, request, jsonify

app = Flask(__name__)

class DummyAuditAgent:
    def __init__(self, agent_id="dummy_audit_agent_v1", api_key=None):
        """
        Initializes the agent using the given or environment-based 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 given smart contract code to the OpenAI API, expecting a JSON-based security analysis.
        Returns a dictionary structured as:
        {
            "agent_id": "some_agent_id",
            "findings": [
                {
                    "finding_id": "...",
                    "severity": "High/Medium/Low",
                    "description": "...",
                    "recommendation": "...",
                    "code_reference": "...",
                },
                ...
            ],
            "metadata": {...}
        }
        """
        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 additional info

Contract code:
{contract_code}

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

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

            raw_text = response.choices[0].message.content.strip()
            parsed_response = json.loads(raw_text)

        except json.JSONDecodeError:
            parsed_response = {
                "agent_id": self.agent_id,
                "findings": [],
                "metadata": {
                    "error": "Failed to parse JSON from the AI response",
                    "raw_response": raw_text
                }
            }
        except Exception as e:
            parsed_response = {
                "agent_id": self.agent_id,
                "findings": [],
                "metadata": {
                    "error": str(e)
                }
            }

        # Ensure the agent_id is set in the final structure
        if isinstance(parsed_response, dict):
            parsed_response.setdefault("agent_id", self.agent_id)

        return parsed_response

# Instantiate the agent (You can override the agent_id or API key if needed)
audit_agent = DummyAuditAgent()

@app.route("/webhook", methods=["POST"])
def webhook():
    """
    Example webhook endpoint that expects a JSON payload containing a `contract_code` field.
    It uses DummyAuditAgent to analyze the code and returns the findings as JSON.
    """
    data = request.get_json(force=True)
    contract_code = data.get("contract_code", "")

    if not contract_code:
        return jsonify({
            "error": "No contract_code provided in the request."
        }), 400

    results = audit_agent.analyze_contract(contract_code)
    return jsonify(results), 200

if __name__ == "__main__":
    # Start the Flask server so the agent can receive webhook requests
    app.run(host="0.0.0.0", port=8000, debug=True)