deepaudits 0.2.0__py3-none-any.whl

deepaudits/__init__.py

@@ -0,0 +1,22 @@
+from .models import (
+    Severity, Finding, AgentOutput, DebateRound, AuditReport,
+    BlindSpot, AnalysisGap, AttackPath, AttackTree, Invariant,
+    EconomicAnalysis, PolicyCheck, QnAResponse,
+)
+from .config import DeepAuditsConfig, load_config, save_global_config
+from .orchestrator import DeepAuditsOrchestrator
+from .memory import AuditMemory
+
+__title__ = "deepaudits"
+__version__ = "0.2.0"
+__author__ = "Nishan Mishra"
+__github__ = "https://github.com/nishanm15"
+__license__ = "MIT"
+
+__all__ = [
+    "Severity", "Finding", "AgentOutput", "DebateRound", "AuditReport",
+    "BlindSpot", "AnalysisGap", "AttackPath", "AttackTree", "Invariant",
+    "EconomicAnalysis", "PolicyCheck", "QnAResponse",
+    "DeepAuditsConfig", "load_config", "save_global_config",
+    "DeepAuditsOrchestrator", "AuditMemory",
+]

deepaudits/__main__.py

@@ -0,0 +1,4 @@
+from deepaudits.cli import app
+
+if __name__ == "__main__":
+    app()

deepaudits/agents/__init__.py

@@ -0,0 +1,13 @@
+from .exploit_hunter import ExploitHunter
+from .solidity_reviewer import SolidityReviewer
+from .economic_attacker import EconomicAttacker
+from .protocol_analyst import ProtocolAnalyst
+from .judge import Judge
+
+__all__ = [
+    "ExploitHunter",
+    "SolidityReviewer",
+    "EconomicAttacker",
+    "ProtocolAnalyst",
+    "Judge"
+]

deepaudits/agents/base_agent.py

@@ -0,0 +1,108 @@
+import re
+import json
+import asyncio
+from typing import Dict, Any, Optional
+from pydantic import ValidationError
+from deepaudits.models import AgentOutput, Finding
+from litellm import acompletion
+
+class BaseAgent:
+    name: str = "BaseAgent"
+    role_description: str = ""
+    prompt_template: str = ""
+
+    def __init__(self, config):
+        self.config = config
+
+    async def analyze(self, contract_code: str, context: dict, verbose: bool = False) -> Optional[AgentOutput]:
+        prompt = self._build_prompt(contract_code, context)
+        
+        for attempt in range(4):
+            try:
+                response_text = await self._call_llm(prompt, verbose)
+                return self._parse_response(response_text)
+            except (json.JSONDecodeError, ValidationError) as e:
+                if attempt == 3:
+                    return None
+                prompt += f"\n\nCRITICAL ERROR: Your previous response caused an error: {str(e)}. You MUST return ONLY valid JSON matching the exact schema."
+            except Exception as e:
+                if attempt == 3:
+                    return None
+                await asyncio.sleep(2)
+        return None
+
+    def _build_prompt(self, contract_code: str, context: dict) -> str:
+        try:
+            with open(self.prompt_template, "r", encoding="utf-8") as f:
+                template = f.read()
+        except FileNotFoundError:
+            template = "System prompt not found."
+        
+        slither_ctx = context.get("slither_summary") or "Not available"
+        ast_ctx = json.dumps(context.get("ast_summary", {}), indent=2)
+        
+        # Multi-pass context injection
+        findings_context = ""
+        if context.get("previous_findings"):
+            findings_context = "\n".join([f"- {f.title} (Severity: {f.severity}, Confidence: {f.confidence})" for f in context["previous_findings"]])
+            findings_context = f"""
+=============================================================
+PREVIOUS FINDINGS (Pass 1)
+=============================================================
+{findings_context}
+"""
+        
+        full_prompt = f"""{template}
+
+=============================================================
+CONTEXT
+=============================================================
+Target Type: {context.get('contract_type', 'Unknown')}
+Profile: {context.get('profile', 'general')}
+
+SLITHER ANALYSIS:
+{slither_ctx}
+
+AST SUMMARY:
+{ast_ctx}
+{findings_context}
+
+=============================================================
+CONTRACT CODE
+=============================================================
+{contract_code}
+"""
+        return full_prompt
+
+    async def _call_llm(self, prompt: str, verbose: bool) -> str:
+        kwargs = {
+            "model": self.config.model,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": self.config.max_tokens_per_agent,
+        }
+        if self.config.api_key:
+            kwargs["api_key"] = self.config.api_key
+        if self.config.api_base:
+            kwargs["api_base"] = self.config.api_base
+
+        try:
+            response = await acompletion(**kwargs)
+            return response.choices[0].message.content
+        except Exception as e:
+            # If the model fails, try a fallback
+            if self.config.model != "gpt-3.5-turbo":
+                print(f"Model {self.config.model} failed, trying fallback...")
+                kwargs["model"] = "gpt-3.5-turbo"
+                response = await acompletion(**kwargs)
+                return response.choices[0].message.content
+            raise
+
+    def _parse_response(self, text: str) -> AgentOutput:
+        json_pattern = re.compile(r'\{.*\}', re.DOTALL)
+        match = json_pattern.search(text)
+        if match:
+            json_str = match.group(0)
+            data = json.loads(json_str)
+            return AgentOutput(**data)
+        else:
+            raise json.JSONDecodeError("No JSON object could be decoded", text, 0)

deepaudits/agents/economic_attacker.py

@@ -0,0 +1,7 @@
+import os
+from .base_agent import BaseAgent
+
+class EconomicAttacker(BaseAgent):
+    name = "EconomicAttacker"
+    role_description = "Specialized DeFi researcher looking for tokenomics and oracle vulnerabilities."
+    prompt_template = os.path.join(os.path.dirname(__file__), "..", "prompts", "economic_attacker.txt")

deepaudits/agents/exploit_hunter.py

@@ -0,0 +1,7 @@
+import os
+from .base_agent import BaseAgent
+
+class ExploitHunter(BaseAgent):
+    name = "ExploitHunter"
+    role_description = "Specialized smart contract security researcher looking for code-level exploits."
+    prompt_template = os.path.join(os.path.dirname(__file__), "..", "prompts", "exploit_hunter.txt")

deepaudits/agents/judge.py

@@ -0,0 +1,105 @@
+import os
+import json
+import asyncio
+from typing import List, Dict, Any
+from deepaudits.models import AgentOutput, DebateRound, AuditReport, Finding
+from deepaudits.agents.base_agent import BaseAgent
+from datetime import datetime
+
+class Judge(BaseAgent):
+    name = "Judge"
+    role_description = "Lead auditor who synthesizes findings and resolves debates."
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.prompt_template = os.path.join(os.path.dirname(__file__), "..", "prompts", "judge.txt")
+
+    async def evaluate(self, agent_outputs: List[AgentOutput], debate_rounds: List[DebateRound], context: dict, verbose: bool = False) -> AuditReport:
+        all_findings_str = ""
+        for out in agent_outputs:
+            all_findings_str += f"\n--- Agent: {out.agent_name} ---\n"
+            for f in out.findings:
+                all_findings_str += json.dumps(f.model_dump(), indent=2) + "\n"
+
+        debate_str = ""
+        for round in debate_rounds:
+            debate_str += f"\nRound {round.round_number}:\n"
+            debate_str += json.dumps(round.model_dump(), indent=2) + "\n"
+
+        try:
+            with open(self.prompt_template, "r", encoding="utf-8") as f:
+                template = f.read()
+        except FileNotFoundError:
+            template = "Judge prompt not found."
+
+        prompt = f"""
+{template}
+
+=============================================================
+ALL AGENT FINDINGS
+=============================================================
+{all_findings_str}
+
+=============================================================
+DEBATE ROUNDS
+=============================================================
+{debate_str}
+"""
+        
+        response_text = ""
+        for attempt in range(4):
+            try:
+                response_text = await self._call_llm(prompt, verbose)
+                return self._parse_judge_response(response_text, agent_outputs, debate_rounds, context)
+            except Exception as e:
+                if attempt == 3:
+                    raise Exception(f"Judge failed after 4 attempts: {e}")
+                prompt += f"\n\nCRITICAL: Your previous response was invalid. Error: {str(e)}. You MUST return ONLY valid JSON."
+                await asyncio.sleep(2)
+        
+        raise Exception("Judge failed")
+
+    def _parse_judge_response(self, text: str, agent_outputs: List[AgentOutput], debate_rounds: List[DebateRound], context: dict) -> AuditReport:
+        import re
+        json_pattern = re.compile(r'\{.*\}', re.DOTALL)
+        match = json_pattern.search(text)
+        if not match:
+            raise json.JSONDecodeError("No JSON object could be decoded", text, 0)
+        
+        data = json.loads(match.group(0))
+        
+        findings = [Finding(**f) for f in data.get("findings", [])]
+        
+        audit_score = 100
+        counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "informational": 0}
+        
+        for f in findings:
+            sev = f.severity.value if hasattr(f.severity, 'value') else f.severity
+            counts[sev] += 1
+            
+        audit_score -= min(counts["critical"] * 25, 50)
+        audit_score -= min(counts["high"] * 10, 30)
+        audit_score -= min(counts["medium"] * 3, 15)
+        audit_score -= min(counts["low"] * 1, 5)
+        audit_score = max(audit_score, 0)
+
+        agents_used = list(set([out.agent_name for out in agent_outputs]))
+
+        return AuditReport(
+            target=context.get("target", "unknown"),
+            timestamp=datetime.utcnow().isoformat(),
+            contract_type=context.get("contract_type", "Unknown"),
+            profile=context.get("profile", "general"),
+            agents_used=agents_used,
+            total_findings=len(findings),
+            critical=counts["critical"],
+            high=counts["high"],
+            medium=counts["medium"],
+            low=counts["low"],
+            informational=counts["informational"],
+            findings=findings,
+            debate_rounds=debate_rounds,
+            slither_output=context.get("slither_summary"),
+            executive_summary=data.get("executive_summary", ""),
+            audit_score=audit_score
+        )

deepaudits/agents/protocol_analyst.py

@@ -0,0 +1,7 @@
+import os
+from .base_agent import BaseAgent
+
+class ProtocolAnalyst(BaseAgent):
+    name = "ProtocolAnalyst"
+    role_description = "System architect looking for logic errors and trust assumption violations."
+    prompt_template = os.path.join(os.path.dirname(__file__), "..", "prompts", "protocol_analyst.txt")

deepaudits/agents/solidity_reviewer.py

@@ -0,0 +1,7 @@
+import os
+from .base_agent import BaseAgent
+
+class SolidityReviewer(BaseAgent):
+    name = "SolidityReviewer"
+    role_description = "Specialized smart contract engineer focusing on code quality and language pitfalls."
+    prompt_template = os.path.join(os.path.dirname(__file__), "..", "prompts", "solidity_reviewer.txt")

deepaudits/analyzer/__init__.py

@@ -0,0 +1,5 @@
+from .preprocessor import Preprocessor
+from .report_builder import ReportBuilder
+from .html_report import HTMLReportBuilder
+
+__all__ = ["Preprocessor", "ReportBuilder", "HTMLReportBuilder"]

deepaudits/analyzer/html_report.py

@@ -0,0 +1,188 @@
+import json
+from pathlib import Path
+from typing import List
+from deepaudits.models import AuditReport
+
+class HTMLReportBuilder:
+    def build(self, report: AuditReport, out_dir: str) -> str:
+        out_path = Path(out_dir)
+        out_path.mkdir(parents=True, exist_ok=True)
+        safe = str(report.target).split("/")[-1].replace("\\", "_").replace(":", "_")
+        ts = report.timestamp.replace(":", "").replace("-", "").replace(".", "_")
+        base = f"deepaudits-{safe}-{ts}"
+        fp = out_path / f"{base}.html"
+        with open(fp, "w") as f:
+            f.write(self._html(report))
+        return str(fp)
+
+    def _html(self, report: AuditReport) -> str:
+        sev_order = {"critical":0,"high":1,"medium":2,"low":3,"informational":4}
+        sorted_f = sorted(report.findings, key=lambda x: (sev_order.get(x.severity.value if hasattr(x.severity,'value') else x.severity,5), -x.confidence))
+
+        sc = "#ff4444" if report.audit_score < 50 else "#ff8800" if report.audit_score < 80 else "#44cc44"
+        colors = {"critical":"#ff4444","high":"#ff8800","medium":"#ffcc00","low":"#4488ff","informational":"#aaa"}
+
+        rows = ""
+        for f in sorted_f:
+            sev = f.severity.value if hasattr(f.severity,'value') else f.severity
+            c = colors.get(sev,"#fff")
+            rows += f"""
+            <tr>
+                <td>{f.id}</td>
+                <td><span style="color:{c};font-weight:bold">{sev.upper()}</span></td>
+                <td>{f.confidence}%</td>
+                <td>{f.title}</td>
+                <td>{f.agent}</td>
+                <td>{f.location}</td>
+                <td><button onclick="tog('{f.id}')">View</button></td>
+            </tr>
+            <tr id="d-{f.id}" style="display:none;background:#1a1a2e">
+                <td colspan="7"><strong>Description:</strong> {f.description}<br>
+                <strong>Impact:</strong> {f.impact}<br>
+                <strong>Exploit Path:</strong> {f.exploit_path}<br>
+                <strong>Recommendation:</strong> {f.recommendation}
+                {f'<br><strong>PoC:</strong><pre>{f.poc_code}</pre>' if f.poc_code else ''}</td>
+            </tr>"""
+
+        # Blind spots
+        bs_rows = ""
+        for bs in report.blind_spots:
+            sev = bs.severity.value if hasattr(bs.severity,'value') else bs.severity
+            c = colors.get(sev,"#fff")
+            bs_rows += f"<tr><td>{bs.id}</td><td style='color:{c}'>{sev.upper()}</td><td>{bs.title}</td><td>{bs.reason}</td><td>{bs.location}</td></tr>"
+
+        # Invariants
+        inv_rows = ""
+        for inv in report.invariants:
+            icon = {"verified":"✅","broken":"❌","unverifiable":"⚠️"}.get(inv.status,"⚪")
+            inv_rows += f"<tr><td>{icon}</td><td>{inv.description}</td><td style='color:{'red' if inv.status=='broken' else 'green' if inv.status=='verified' else '#ffcc00'}'>{inv.status.upper()}</td></tr>"
+
+        # Economic
+        econ_rows = ""
+        for ea in report.economic_analyses:
+            vc = "red" if "not profitable" in ea.verdict.lower() else "green" if "highly profitable" in ea.verdict.lower() else "#ffcc00"
+            econ_rows += f"<tr><td>{ea.finding_id}</td><td>{ea.capital_required}</td><td>{ea.gas_cost}</td><td>{ea.potential_profit}</td><td>{ea.risk_reward_ratio}</td><td style='color:{vc}'>{ea.verdict}</td></tr>"
+
+        # Attack tree
+        at_html = ""
+        for tree in report.attack_tree:
+            at_html += f"<h3>🎯 {tree.goal}</h3>"
+            for path in tree.paths:
+                icon = {"confirmed":"🟢","suspected":"🟡","blocked":"🔴"}.get(path.status,"⚪")
+                at_html += f"<div style='margin:8px 0;padding:8px;background:#1a1a2e;border-radius:4px'>{icon} <b>{path.goal}</b> ({path.status})"
+                for i, s in enumerate(path.steps, 1):
+                    at_html += f"<br>&nbsp;&nbsp;{i}. {s}"
+                at_html += "</div>"
+
+        # Policy checks
+        pc_rows = ""
+        for pc in report.policy_checks:
+            icon = {"passed":"✅","violated":"❌","unverifiable":"⚠️"}.get(pc.status,"⚪")
+            color = "green" if pc.status=="passed" else "red" if pc.status=="violated" else "#ffcc00"
+            pc_rows += f"<tr><td>{icon}</td><td>{pc.policy_text}</td><td style='color:{color}'>{pc.status.upper()}</td><td>{pc.details}</td></tr>"
+
+        html = f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>DeepAudits Report</title>
+<style>
+*{{margin:0;padding:0;box-sizing:border-box}}
+body{{font-family:'Segoe UI',sans-serif;background:#0a0a1a;color:#e0e0e0;padding:20px}}
+.container{{max-width:1300px;margin:0 auto}}
+h1{{color:#00d4ff;margin-bottom:10px}}
+h2{{color:#00d4ff;margin:25px 0 10px;border-bottom:1px solid #333;padding-bottom:5px}}
+.score-box{{display:inline-block;padding:20px 40px;border-radius:10px;font-size:48px;font-weight:bold;text-align:center;margin:20px 0}}
+.score-box .label{{font-size:14px;color:#aaa}}
+.summary-grid{{display:grid;grid-template-columns:repeat(auto-fit,minmax(120px,1fr));gap:10px;margin:15px 0}}
+.summary-item{{background:#1a1a2e;padding:15px;border-radius:8px;text-align:center}}
+.summary-item .count{{font-size:28px;font-weight:bold}}
+.summary-item .label{{font-size:12px;color:#888}}
+table{{width:100%;border-collapse:collapse;margin:15px 0;background:#1a1a2e;border-radius:8px;overflow:hidden}}
+th,td{{padding:10px 15px;text-align:left;border-bottom:1px solid #333}}
+th{{background:#16213e;color:#00d4ff;font-weight:600}}
+tr:hover{{background:#1a1a3e}}
+button{{background:#00d4ff;color:#000;border:none;padding:5px 10px;border-radius:4px;cursor:pointer;font-weight:bold}}
+button:hover{{background:#00b8e6}}
+pre{{background:#111;padding:10px;border-radius:4px;overflow-x:auto;margin:8px 0}}
+.exec-summary{{background:#1a1a2e;padding:20px;border-radius:8px;line-height:1.6;margin:15px 0}}
+.meta{{color:#888;font-size:14px}}
+.tab{{overflow:hidden;margin:15px 0}}
+.tab button{{background:#16213e;color:#00d4ff;float:left;border:none;outline:none;cursor:pointer;padding:12px 20px;transition:0.3s;border-radius:4px 4px 0 0;margin-right:2px}}
+.tab button:hover{{background:#1e2d50}}
+.tab button.active{{background:#00d4ff;color:#000}}
+.tabcontent{{display:none;padding:20px;background:#1a1a2e;border-radius:0 4px 4px 4px}}
+@media(max-width:768px){{th,td{{font-size:12px;padding:6px}}}}
+</style>
+</head>
+<body>
+<div class="container">
+<h1>🔍 DeepAudits Security Report</h1>
+<p class="meta"><strong>Target:</strong> {report.target} | <strong>Score:</strong> {report.audit_score}/100 | <strong>{report.contract_type}</strong></p>
+
+<div style="text-align:center">
+<div class="score-box" style="background:{sc}22;border:3px solid {sc}">
+<span style="color:{sc}">{report.audit_score}</span>
+<div class="label">SECURITY SCORE</div></div></div>
+
+<h2>📊 Findings</h2>
+<div class="summary-grid">
+<div class="summary-item"><div class="count" style="color:#ff4444">{report.critical}</div><div class="label">Critical</div></div>
+<div class="summary-item"><div class="count" style="color:#ff8800">{report.high}</div><div class="label">High</div></div>
+<div class="summary-item"><div class="count" style="color:#ffcc00">{report.medium}</div><div class="label">Medium</div></div>
+<div class="summary-item"><div class="count" style="color:#4488ff">{report.low}</div><div class="label">Low</div></div>
+<div class="summary-item"><div class="count" style="color:#aaa">{report.informational}</div><div class="label">Info</div></div>
+<div class="summary-item"><div class="count" style="color:#00d4ff">{report.total_findings}</div><div class="label">Total</div></div>
+</div>
+
+<h2>📝 Executive Summary</h2>
+<div class="exec-summary">{report.executive_summary}</div>
+
+<div class="tab">
+<button class="tablinks active" onclick="openTab(event,'findings')">🔎 Findings {report.total_findings}</button>
+<button class="tablinks" onclick="openTab(event,'blindspots')">⚠️ Blind Spots {len(report.blind_spots)}</button>
+<button class="tablinks" onclick="openTab(event,'invariants')">🧮 Invariants {len(report.invariants)}</button>
+<button class="tablinks" onclick="openTab(event,'economic')">💰 Economics {len(report.economic_analyses)}</button>
+<button class="tablinks" onclick="openTab(event,'attacktree')">🗺️ Attack Tree</button>
+{"<button class='tablinks' onclick=\"openTab(event,'policies')\">📋 Policies</button>" if report.policy_checks else ""}
+</div>
+
+<div id="findings" class="tabcontent" style="display:block">
+<table><thead><tr><th>ID</th><th>Severity</th><th>Conf</th><th>Title</th><th>Agent</th><th>Location</th><th></th></tr></thead>
+<tbody>{rows}</tbody></table></div>
+
+<div id="blindspots" class="tabcontent">
+{"<table><tr><th>ID</th><th>Sev</th><th>Title</th><th>Reason</th><th>Location</th></tr>"+bs_rows+"</table>" if report.blind_spots else "<p>No blind spots identified.</p>"}
+{("<h3>Analysis Gaps</h3><ul>"+''.join(f"<li><b>{g.title}</b> ({g.location}): {g.description}</li>" for g in report.analysis_gaps)+"</ul>") if report.analysis_gaps else ""}
+</div>
+
+<div id="invariants" class="tabcontent">
+{"<table><tr><th></th><th>Invariant</th><th>Status</th></tr>"+inv_rows+"</table>" if report.invariants else "<p>No invariants inferred.</p>"}
+</div>
+
+<div id="economic" class="tabcontent">
+{"<table><tr><th>ID</th><th>Capital</th><th>Gas</th><th>Profit</th><th>Risk/Reward</th><th>Verdict</th></tr>"+econ_rows+"</table>" if report.economic_analyses else "<p>No economic analysis performed.</p>"}
+</div>
+
+<div id="attacktree" class="tabcontent">
+{at_html if report.attack_tree else "<p>No attack trees generated.</p>"}
+</div>
+
+{"<div id='policies' class='tabcontent'><table><tr><th></th><th>Policy</th><th>Status</th><th>Details</th></tr>"+pc_rows+"</table></div>" if report.policy_checks else ""}
+
+<h2>🛠 Agents Used</h2>
+<p>{", ".join(report.agents_used)}</p>
+
+<h2>⚙️ Slither</h2>
+<pre>{report.slither_output or "Skipped."}</pre>
+
+<p class="meta" style="margin-top:30px;text-align:center">DeepAudits v0.2.0</p>
+</div>
+<script>
+function tog(id){{var e=document.getElementById('d-'+id);e.style.display=e.style.display==='none'?'table-row':'none'}}
+function openTab(e,t){{var i,x=document.getElementsByClassName('tabcontent');for(i=0;i<x.length;i++)x[i].style.display='none';var b=document.getElementsByClassName('tablinks');for(i=0;i<b.length;i++)b[i].classList.remove('active');document.getElementById(t).style.display='block';e.currentTarget.classList.add('active')}}
+</script>
+</body>
+</html>"""
+        return html

deepaudits/analyzer/preprocessor.py

@@ -0,0 +1,136 @@
+import re
+import hashlib
+from pathlib import Path
+from typing import List, Tuple, Dict
+
+class Preprocessor:
+    def __init__(self, cache_dir: str = ".deepaudits/cache"):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def _compute_hash(self, source_code: str) -> str:
+        """Compute hash of source code for caching."""
+        return hashlib.md5(source_code.encode()).hexdigest()
+
+    def _get_cache_path(self, file_hash: str) -> Path:
+        return self.cache_dir / f"ast_{file_hash}.json"
+
+    def chunk_contract(self, source_code: str, max_lines: int = 2000, overlap: int = 200) -> List[Dict[str, str]]:
+        """Smart chunking: chunk by functions, not just lines."""
+        # First, try to extract functions
+        functions = self._extract_functions_with_context(source_code)
+        if functions:
+            return functions
+        
+        # Fallback: line-based chunking
+        lines = source_code.splitlines()
+        if len(lines) <= max_lines:
+            return [{"type": "code", "content": source_code}]
+        
+        chunks = []
+        i = 0
+        while i < len(lines):
+            chunk_lines = lines[i:i + max_lines]
+            chunks.append({"type": "code", "content": "\n".join(chunk_lines)})
+            i += (max_lines - overlap)
+        return chunks
+
+    def _extract_functions_with_context(self, source_code: str) -> List[Dict[str, str]]:
+        """Extract functions with state variable context for smart chunking."""
+        lines = source_code.splitlines()
+        
+        # Extract state variables from the contract
+        state_vars = self._extract_state_variables(source_code)
+        
+        # Extract all function definitions with their full body
+        function_pattern = re.compile(
+            r'(//.*?\n)?'                          # Optional comment before function
+            r'(function\s+([a-zA-Z0-9_]+)\s*\((.*?)\)\s*'  # Function signature
+            r'((?:public|external|internal|private|view|pure|payable|override|virtual|abstract|\s)*)' # Modifiers
+            r'\s*(?:returns\s*\(.*?\))?\s*\{',     # Optional return and opening brace
+            re.MULTILINE | re.DOTALL
+        )
+        
+        matches = list(function_pattern.finditer(source_code))
+        if not matches:
+            return []
+        
+        chunks = []
+        for match in matches:
+            func_start = match.start()
+            func_signature = match.group(0)
+            
+            # Find the matching closing brace
+            brace_count = 1
+            pos = match.end()
+            while pos < len(source_code) and brace_count > 0:
+                if source_code[pos] == '{':
+                    brace_count += 1
+                elif source_code[pos] == '}':
+                    brace_count -= 1
+                pos += 1
+            
+            func_body = source_code[func_start:pos]
+            
+            # Build context: state variables that are used in this function
+            used_vars = self._get_used_state_vars(func_body, state_vars)
+            
+            # Create chunk with function body and relevant context
+            chunk_content = self._build_chunk(func_signature, func_body, used_vars, state_vars)
+            
+            chunks.append({
+                "type": "function",
+                "function_name": match.group(3),
+                "content": chunk_content
+            })
+        
+        return chunks
+
+    def _extract_state_variables(self, source_code: str) -> Dict[str, str]:
+        """Extract state variables with their types."""
+        state_vars = {}
+        
+        # Pattern for state variable declarations
+        # uint256 public balance;
+        # mapping(address => uint256) public balances;
+        # address private owner;
+        var_patterns = [
+            # Simple type declaration
+            r'(uint[0-9]*|int[0-9]*|address|bool|string|bytes[0-9]*|mapping\([^)]+\))\s+(?:public|private|internal)?\s*([a-zA-Z0-9_]+)\s*[;=]',
+            # Array declaration
+            r'(uint[0-9]*|int[0-9]*|address|bool|string|bytes[0-9]*)\[\]\s+(?:public|private|internal)?\s*([a-zA-Z0-9_]+)\s*[;=]',
+        ]
+        
+        for pattern in var_patterns:
+            for match in re.finditer(pattern, source_code):
+                var_type = match.group(1)
+                var_name = match.group(2)
+                state_vars[var_name] = var_type
+        
+        return state_vars
+
+    def _get_used_state_vars(self, func_body: str, state_vars: Dict[str, str]) -> List[str]:
+        """Find which state variables are used in a function."""
+        used = []
+        for var_name in state_vars:
+            # Check if variable is referenced in the function body
+            # Use word boundary to avoid partial matches
+            if re.search(r'\b' + re.escape(var_name) + r'\b', func_body):
+                used.append(var_name)
+        return used
+
+    def _build_chunk(self, func_signature: str, func_body: str, used_vars: List[str], all_vars: Dict[str, str]) -> str:
+        """Build a chunk with function and relevant state variable context."""
+        context_parts = []
+        
+        # Add relevant state variables
+        if used_vars:
+            context_parts.append("// Relevant State Variables:\n")
+            for var in used_vars:
+                if var in all_vars:
+                    context_parts.append(f"// {all_vars[var]} {var};\n")
+        
+        # Add the function body
+        context_parts.append(func_body)
+        
+        return "".join(context_parts)