File size: 6,885 Bytes
e545bf5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
"""Optional AI assistance for Calgary SWMR drafting and consistency review.

This module is intentionally separate from the general model-analysis agent.
The deterministic report engine remains the source of all engineering values,
criteria statuses, checklist statuses, and permitted conclusions.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Mapping, Sequence

import requests

PROMPT_DIR = Path(__file__).resolve().parent / "prompts"

PROVIDERS: dict[str, dict[str, Any]] = {
    "Claude (Anthropic)": {
        "models": ["claude-opus-4-5", "claude-sonnet-4-5", "claude-haiku-4-5"],
        "format": "anthropic",
        "url": "https://api.anthropic.com/v1/messages",
    },
    "GPT (OpenAI)": {
        "models": ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo"],
        "format": "openai",
        "url": "https://api.openai.com/v1/chat/completions",
    },
    "Gemini (Google)": {
        "models": ["gemini-2.0-flash", "gemini-1.5-pro", "gemini-1.5-flash"],
        "format": "gemini",
        "url": "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
    },
    "Groq (Free tier)": {
        "models": ["llama-3.3-70b-versatile", "llama-3.1-8b-instant", "mixtral-8x7b-32768"],
        "format": "openai",
        "url": "https://api.groq.com/openai/v1/chat/completions",
    },
    "Mistral": {
        "models": ["mistral-large-latest", "mistral-small-latest", "open-mistral-7b"],
        "format": "openai",
        "url": "https://api.mistral.ai/v1/chat/completions",
    },
}


def load_prompt(name: str) -> str:
    path = PROMPT_DIR / name
    if not path.exists():
        raise FileNotFoundError(f"Prompt file not found: {path}")
    return path.read_text(encoding="utf-8").strip()


DRAFT_SYSTEM_PROMPT = load_prompt("calgary_report_drafting_system.txt")
REVIEW_SYSTEM_PROMPT = load_prompt("calgary_report_review_system.txt")


def _call_provider(
    provider_name: str,
    api_key: str,
    model: str,
    messages: Sequence[Mapping[str, str]],
    system_prompt: str,
    max_tokens: int = 3200,
    timeout: int = 90,
) -> str:
    if provider_name not in PROVIDERS:
        raise ValueError(f"Unsupported provider: {provider_name}")
    if not api_key.strip():
        raise ValueError("An API key is required for optional AI assistance.")

    provider = PROVIDERS[provider_name]
    fmt = provider["format"]

    if fmt == "anthropic":
        response = requests.post(
            provider["url"],
            headers={
                "Content-Type": "application/json",
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
            },
            json={
                "model": model,
                "max_tokens": max_tokens,
                "system": system_prompt,
                "messages": list(messages),
            },
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()["content"][0]["text"].strip()

    if fmt == "openai":
        response = requests.post(
            provider["url"],
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
            json={
                "model": model,
                "max_tokens": max_tokens,
                "messages": [{"role": "system", "content": system_prompt}] + list(messages),
            },
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"].strip()

    if fmt == "gemini":
        url = provider["url"].replace("{model}", model) + f"?key={api_key}"
        contents = []
        for message in messages:
            role = "model" if message["role"] == "assistant" else "user"
            contents.append({"role": role, "parts": [{"text": message["content"]}]})
        response = requests.post(
            url,
            json={
                "system_instruction": {"parts": [{"text": system_prompt}]},
                "contents": contents,
                "generationConfig": {"maxOutputTokens": max_tokens},
            },
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()["candidates"][0]["content"]["parts"][0]["text"].strip()

    raise ValueError(f"Provider format is not implemented: {fmt}")


def _safe_json(data: Mapping[str, Any]) -> str:
    return json.dumps(data, ensure_ascii=False, default=str, separators=(",", ":"))


def draft_report_section(
    *,
    provider_name: str,
    api_key: str,
    model: str,
    section: str,
    report_context: Mapping[str, Any],
    additional_instruction: str = "",
) -> str:
    """Draft narrative from deterministic report context only."""
    request = f"""Prepare this Calgary SWMR draft component: {section}.

Use only the VERIFIED_REPORT_CONTEXT JSON below. Follow every conclusion-control,
checklist, storage, outfall, minor-system, major-system, and model-input/output rule
in the system prompt. Preserve all values and units exactly.

VERIFIED_REPORT_CONTEXT:
{_safe_json(report_context)}

ADDITIONAL_USER_INSTRUCTION:
{additional_instruction.strip() or 'None'}

Return report-ready prose with clear headings. Do not include a preamble about being an AI.
"""
    return _call_provider(
        provider_name,
        api_key,
        model,
        [{"role": "user", "content": request}],
        DRAFT_SYSTEM_PROMPT,
    )


def review_report_narrative(
    *,
    provider_name: str,
    api_key: str,
    model: str,
    narrative: str,
    report_context: Mapping[str, Any],
) -> str:
    """Review narrative against deterministic facts without changing model results."""
    request = f"""Review the DRAFT_NARRATIVE against VERIFIED_REPORT_CONTEXT.
Identify unsupported claims, value or unit discrepancies, checklist omissions,
misuse of criteria, overstatements, and contradictions. Then provide corrected
replacement wording for each material issue.

VERIFIED_REPORT_CONTEXT:
{_safe_json(report_context)}

DRAFT_NARRATIVE:
{narrative}
"""
    return _call_provider(
        provider_name,
        api_key,
        model,
        [{"role": "user", "content": request}],
        REVIEW_SYSTEM_PROMPT,
    )


def draft_multiple_report_sections(*, provider_name: str, api_key: str, model: str, sections: Sequence[str], report_context: Mapping[str, Any], additional_instruction: str = "") -> dict[str, str]:
    """Generate independent section drafts so each can be reviewed and approved."""
    drafts: dict[str, str] = {}
    for section in sections:
        drafts[section] = draft_report_section(
            provider_name=provider_name, api_key=api_key, model=model, section=section,
            report_context=report_context, additional_instruction=additional_instruction,
        )
    return drafts