"""
Legal draft generation core.
Two parallel implementations (GPT-4o and Claude) sit behind a single
interface, `DraftGenerator`. A `DraftSession` wraps a generator and
keeps conversation history + refinement count, so stateful UIs only
need a few lines of glue.
Stateless surface (good for Lambda / FastAPI):
gen = get_generator("claude") # or "gpt"
result = gen.generate_initial(...) # returns InitialDraftResult
result = gen.refine(message, conversation_history, user_context)
Stateful convenience (good for Streamlit / Gradio):
session = DraftSession(model="claude")
result = session.generate_initial(...)
result = session.refine("add more precedents")
session.draft_markdown # latest version
session.remaining_refinements
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Literal
from .clients import get_anthropic_client, get_openai_client
from .config import get_settings
from .retrieval import format_citations_markdown, search_for_draft
from .templates import parse_json_array
MAX_REFINEMENTS = 5
Message = dict # {"role": "user" | "assistant", "content": str}
# ===================================================================
# PROMPTS
# ===================================================================
SYSTEM_PROMPT_DRAFT_CLAUDE = """You are an expert legal drafter specialising in Pakistani law.
Hard rules:
- Output ONLY the petition in Markdown. No preamble, no closing remarks, no meta-commentary.
- Never assume facts not given. Use [PLACEHOLDER] where information is missing.
- Follow the structure exactly as instructed in the user turn.
- Integrate every cited law or precedent given to you; do not invent citations.
- Maintain formal, court-ready language throughout."""
SYSTEM_PROMPT_DRAFT_GPT = """You are an expert legal drafter for Pakistani law. Your task is to create a professional, court-ready legal petition in MARKDOWN format using the provided inputs:
1. User Input: Case details including client info, petition type, court, facts, relevant laws, and sections.
2. Knowledge Base Context: Relevant laws, case precedents, and ordinances.
3. Template Draft Analysis: Structure and style guidelines.
Instructions:
- Create a complete, professional legal draft
- Follow the template structure and style
- Integrate legal context with accurate citations
- Use MARKDOWN format with proper headings
- Be comprehensive and court-ready
- ALWAYS adhere to the user's custom instructions/preferences if provided
- Output ONLY the draft, no explanations"""
SYSTEM_PROMPT_REFINE = """You are an expert Pakistani legal drafter continuing to refine a petition.
Hard rules:
- Apply ONLY the changes the user requests; preserve everything else exactly.
- Output the COMPLETE updated petition in Markdown. No explanations before or after.
- Never invent facts or citations.
- Always follow any custom instructions the user has provided."""
SYSTEM_PROMPT_QUERY_CLAUDE = """You are a Pakistani legal research assistant.
Given a case description, produce a JSON array of concise search queries (strings only).
The queries will be run against a vector database containing:
- The Constitution of Pakistan 1973
- Punjab High Court and Supreme Court case law
- FBR Income Tax Ordinance 2001 and Sales Tax Act 1990
Rules:
- Return ONLY a valid JSON array of strings. No markdown fences, no explanation.
- Each query should target a specific legal concept, statute, or precedent relevant to the case.
- Generate between 8 and 15 queries, ordered from most to least specific."""
SYSTEM_PROMPT_QUERY_GPT = (
"You are a legal research assistant. "
"A new petition needs to be drafted using the following client/case description. "
"Devise 5-6 or more concise queries that will be helpful to retrieve relevant information "
"from a knowledge base containing the Constitution of Pakistan, Punjab case law, "
"and FBR tax ordinances. "
"Return ONLY a JSON array of strings, no extra text."
)
# ===================================================================
# RESULT TYPES
# ===================================================================
@dataclass
class InitialDraftResult:
draft_markdown: str # the petition only, no references appended
draft_with_citations: str # same draft with `### References` appended (if requested)
citations: list[dict] # [{"score": float, "source": str}, ...]
context_text: str # retrieved KB passages
template_analysis: str # the template style guide that was used
conversation_history: list[Message] # carry this into refine()
@dataclass
class RefinementResult:
draft_markdown: str
draft_with_citations: str
conversation_history: list[Message]
refinement_number: int # 1..MAX_REFINEMENTS
remaining: int # MAX_REFINEMENTS - refinement_number
# ===================================================================
# PROMPT BUILDERS (shared)
# ===================================================================
def _build_user_prompt(
case_text: str,
context_text: str,
template_analysis: str,
user_context: str,
) -> str:
"""Build the rich user-turn that both models consume for the initial draft."""
custom_section = ""
if user_context and user_context.strip():
custom_section = (
"\n\n\n"
"The user has provided the following persistent preferences. "
"You MUST follow them throughout this draft and all subsequent refinements:\n"
f"{user_context.strip()}\n"
""
)
return f"""
Draft a complete, court-ready legal petition in Markdown for a Pakistani court.
{custom_section}
{case_text}
The following passages are retrieved from a vector database of Pakistani law and case precedents.
Integrate the most relevant ones as citations in the appropriate sections of the petition.
{context_text or "(no matches found)"}
Follow the structure, tone, length, and formatting described below for the petition.
{template_analysis or "(no template - use standard Pakistani High Court petition format)"}
1. Title
2. Parties (Petitioner / Respondent)
3. Jurisdiction
4. Statement of Facts
5. Grounds for Petition
6. Legal Arguments (with citations from knowledge base)
7. Prayer / Relief Sought
8. Interim Relief (if applicable)
9. Verification / Affidavit
Output ONLY the petition in Markdown. No preamble or closing remarks."""
def _build_refinement_prompt(message: str, user_context: str) -> str:
context_section = ""
if user_context and user_context.strip():
context_section = (
"\n\n**User's persistent custom instructions - always follow these:**\n"
f"{user_context.strip()}\n"
)
return (
f"The user wants the following change to the petition:{context_section}\n\n"
f"**Request:** {message}\n\n"
"Apply the change and return the COMPLETE updated petition in Markdown."
)
# ===================================================================
# ABSTRACT BASE
# ===================================================================
class DraftGenerator(ABC):
"""Common interface for any LLM-backed draft generator."""
name: str = "abstract"
@abstractmethod
def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
...
@abstractmethod
def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
"""Return (draft_markdown, initial_conversation_history)."""
@abstractmethod
def _refine_call(self, conversation_history: list[Message]) -> str:
"""Return updated draft markdown."""
# ---------- public API ----------
def generate_initial(
self,
case_text: str,
*,
user_context: str = "",
template_analysis: str = "",
add_citations: bool = True,
retrieval_top_k: int = 10,
retrieval_max_chars: int = 10_000,
) -> InitialDraftResult:
if not case_text or not case_text.strip():
raise ValueError("case_text is required")
# 1. Knowledge base retrieval
queries = self.build_queries(case_text)
context_text, citations = search_for_draft(
queries,
top_k=retrieval_top_k,
max_chars=retrieval_max_chars,
)
# 2. Build the user-turn prompt
user_prompt = _build_user_prompt(
case_text=case_text,
context_text=context_text,
template_analysis=template_analysis,
user_context=user_context,
)
# 3. Call the model
draft_md, conversation_history = self._initial_call(user_prompt)
# 4. Optionally append citations for display / docx
draft_with_citations = draft_md
if add_citations:
draft_with_citations += format_citations_markdown(citations)
return InitialDraftResult(
draft_markdown=draft_md,
draft_with_citations=draft_with_citations,
citations=citations,
context_text=context_text,
template_analysis=template_analysis,
conversation_history=conversation_history,
)
def refine(
self,
message: str,
conversation_history: list[Message],
*,
user_context: str = "",
citations: list[dict] | None = None,
add_citations: bool = True,
refinement_number: int = 1,
) -> RefinementResult:
if not message or not message.strip():
raise ValueError("message is required")
if refinement_number > MAX_REFINEMENTS:
raise ValueError(
f"Refinement limit ({MAX_REFINEMENTS}) reached."
)
instruction = _build_refinement_prompt(message, user_context)
new_history = list(conversation_history) + [
{"role": "user", "content": instruction}
]
draft_md = self._refine_call(new_history)
new_history.append({"role": "assistant", "content": draft_md})
draft_with_citations = draft_md
if add_citations and citations:
draft_with_citations += format_citations_markdown(citations)
return RefinementResult(
draft_markdown=draft_md,
draft_with_citations=draft_with_citations,
conversation_history=new_history,
refinement_number=refinement_number,
remaining=MAX_REFINEMENTS - refinement_number,
)
# ===================================================================
# CLAUDE IMPLEMENTATION
# ===================================================================
class ClaudeDraftGenerator(DraftGenerator):
name = "claude"
def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
try:
resp = get_anthropic_client().messages.create(
model=get_settings().claude_model,
max_tokens=600,
system=SYSTEM_PROMPT_QUERY_CLAUDE,
messages=[{"role": "user", "content": case_text}],
)
return parse_json_array(resp.content[0].text.strip(), case_text)[:max_queries]
except Exception:
return [case_text[:512]]
def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
resp = get_anthropic_client().messages.create(
model=get_settings().claude_model,
max_tokens=15_000,
system=SYSTEM_PROMPT_DRAFT_CLAUDE,
messages=[{"role": "user", "content": user_prompt}],
)
draft_md = resp.content[0].text.strip()
history = [
{"role": "user", "content": user_prompt},
{"role": "assistant", "content": draft_md},
]
return draft_md, history
def _refine_call(self, conversation_history: list[Message]) -> str:
resp = get_anthropic_client().messages.create(
model=get_settings().claude_model,
max_tokens=15_000,
system=SYSTEM_PROMPT_REFINE,
messages=conversation_history,
)
return resp.content[0].text.strip()
# ===================================================================
# GPT IMPLEMENTATION
# ===================================================================
class GPTDraftGenerator(DraftGenerator):
name = "gpt"
def build_queries(self, case_text: str, max_queries: int = 15) -> list[str]:
try:
resp = get_openai_client().chat.completions.create(
model=get_settings().gpt_model,
messages=[
{"role": "system", "content": SYSTEM_PROMPT_QUERY_GPT},
{"role": "user", "content": case_text},
],
temperature=0.1,
max_tokens=2000,
)
return parse_json_array(resp.choices[0].message.content.strip(), case_text)[:max_queries]
except Exception:
return [case_text[:512]]
def _initial_call(self, user_prompt: str) -> tuple[str, list[Message]]:
# The GPT version keeps the system message in the conversation history
# so that subsequent refinement turns retain the same instructions.
history = [
{"role": "system", "content": SYSTEM_PROMPT_DRAFT_GPT},
{"role": "user", "content": user_prompt},
]
resp = get_openai_client().chat.completions.create(
model=get_settings().gpt_draft_model,
messages=history,
max_tokens=15_000,
)
draft_md = resp.choices[0].message.content.strip()
history.append({"role": "assistant", "content": draft_md})
return draft_md, history
def _refine_call(self, conversation_history: list[Message]) -> str:
resp = get_openai_client().chat.completions.create(
model=get_settings().gpt_draft_model,
messages=conversation_history,
max_tokens=15_000,
)
return resp.choices[0].message.content.strip()
# ===================================================================
# FACTORY
# ===================================================================
ModelName = Literal["gpt", "claude"]
def get_generator(model: ModelName) -> DraftGenerator:
if model == "claude":
return ClaudeDraftGenerator()
if model == "gpt":
return GPTDraftGenerator()
raise ValueError(f"Unknown model: {model!r}")
# ===================================================================
# STATEFUL CONVENIENCE WRAPPER
# ===================================================================
@dataclass
class DraftSession:
"""
Holds session state for a single draft conversation.
Usage:
s = DraftSession(model="claude", user_context="Always cite the Constitution.")
s.generate_initial(case_text="...", template_analysis="...")
s.refine("Add more precedents")
s.refine("Shorten the facts section")
print(s.draft_with_citations)
"""
model: ModelName = "claude"
user_context: str = ""
add_citations: bool = True
# Populated by generate_initial:
case_text: str = ""
citations: list[dict] = field(default_factory=list)
context_text: str = ""
template_analysis: str = ""
draft_markdown: str = ""
draft_with_citations: str = ""
conversation_history: list[Message] = field(default_factory=list)
user_message_count: int = 0
def __post_init__(self):
self._generator = get_generator(self.model)
# ---------- properties ----------
@property
def remaining_refinements(self) -> int:
return MAX_REFINEMENTS - self.user_message_count
@property
def at_refinement_limit(self) -> bool:
return self.user_message_count >= MAX_REFINEMENTS
# ---------- actions ----------
def generate_initial(
self,
case_text: str,
*,
template_analysis: str = "",
) -> InitialDraftResult:
result = self._generator.generate_initial(
case_text=case_text,
user_context=self.user_context,
template_analysis=template_analysis,
add_citations=self.add_citations,
)
self.case_text = case_text
self.template_analysis = template_analysis
self.citations = result.citations
self.context_text = result.context_text
self.draft_markdown = result.draft_markdown
self.draft_with_citations = result.draft_with_citations
self.conversation_history = result.conversation_history
self.user_message_count = 0
return result
def refine(self, message: str) -> RefinementResult:
if self.at_refinement_limit:
raise RuntimeError(
f"Refinement limit ({MAX_REFINEMENTS}) reached. "
"Start a new session to continue."
)
self.user_message_count += 1
result = self._generator.refine(
message=message,
conversation_history=self.conversation_history,
user_context=self.user_context,
citations=self.citations,
add_citations=self.add_citations,
refinement_number=self.user_message_count,
)
self.draft_markdown = result.draft_markdown
self.draft_with_citations = result.draft_with_citations
self.conversation_history = result.conversation_history
return result
# ---------- serialisation (useful for FastAPI / DynamoDB) ----------
def to_dict(self) -> dict:
return {
"model": self.model,
"user_context": self.user_context,
"add_citations": self.add_citations,
"case_text": self.case_text,
"citations": self.citations,
"context_text": self.context_text,
"template_analysis": self.template_analysis,
"draft_markdown": self.draft_markdown,
"draft_with_citations": self.draft_with_citations,
"conversation_history": self.conversation_history,
"user_message_count": self.user_message_count,
}
@classmethod
def from_dict(cls, data: dict) -> "DraftSession":
s = cls(
model=data["model"],
user_context=data.get("user_context", ""),
add_citations=data.get("add_citations", True),
)
s.case_text = data.get("case_text", "")
s.citations = data.get("citations", [])
s.context_text = data.get("context_text", "")
s.template_analysis = data.get("template_analysis", "")
s.draft_markdown = data.get("draft_markdown", "")
s.draft_with_citations = data.get("draft_with_citations", "")
s.conversation_history = data.get("conversation_history", [])
s.user_message_count = data.get("user_message_count", 0)
return s