How to Ingest a Legal Contract with AutoGen and EnConvert

import asyncio
import json
import os
import time
from typing import Annotated

import requests
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient

ENCONVERT_API_KEY = os.environ["ENCONVERT_API_KEY"]
HEADERS = {"X-API-Key": ENCONVERT_API_KEY}

# Step 1: Upload the legal contract to EnConvert Ingest (multipart, field name "files")
with open("/path/to/your/master-services-agreement.pdf", "rb") as contract_file:
    upload_response = requests.post(
        "https://api.enconvert.com/v2/ingest/files",
        headers=HEADERS,
        files={
            "files": (
                "master-services-agreement.pdf",
                contract_file,
                "application/pdf",
            )
        },
        timeout=60,
    )
upload_response.raise_for_status()
job_id = upload_response.json()["job_id"]

# Step 2: Poll until the job completes
while True:
    job_response = requests.get(
        f"https://api.enconvert.com/v2/ingest/{job_id}",
        headers=HEADERS,
        timeout=60,
    )
    job_response.raise_for_status()
    job = job_response.json()
    if job["status"] == "completed":
        break
    if job["status"] in ("failed", "canceled"):
        raise RuntimeError(
            f"Ingest job {job_id} {job['status']}: {job.get('error_message')}"
        )
    time.sleep(2)

# Step 3: Download the JSONL output, one chunk per line.
# The signed output_url (valid 15 minutes) is fetched with a plain GET, no API key.
jsonl_response = requests.get(job["output_url"], timeout=60)
jsonl_response.raise_for_status()
chunks = [
    json.loads(line)
    for line in jsonl_response.content.splitlines()
    if line.strip()
]
print(f"Loaded {len(chunks)} chunks ready for retrieval")


# Step 4: Define the retrieval tool. AutoGen builds the tool schema from the
# signature, the Annotated descriptions and the docstring.
def retrieve_clauses(
    query: Annotated[str, "Topic to look up, such as termination or liability cap"],
    section: Annotated[str, "Optional heading breadcrumb substring to filter by"] = "",
) -> str:
    """Retrieve the contract sections most relevant to a topic, labelled by section and heading path."""
    terms = set(query.lower().split())
    scored = []

    for chunk in chunks:
        metadata = chunk["metadata"]
        # headings_path is a list of headings, e.g. ["Termination", "For Cause"]
        breadcrumb = " > ".join(metadata["headings_path"])
        section_name = metadata["section"]

        if section and section.lower() not in f"{breadcrumb} {section_name}".lower():
            continue

        searchable = f"{breadcrumb} {chunk['content']}".lower()
        score = sum(term in searchable for term in terms)
        if score or section:
            label = f"Section: {section_name} | Path: {breadcrumb}"
            scored.append((score, label, chunk["content"]))

    scored.sort(key=lambda item: item[0], reverse=True)

    results = [f"[{label}]\n{content}" for _, label, content in scored[:4]]
    return "\n\n---\n\n".join(results) if results else "No matching clauses found."


# Replace keyword scoring with vector similarity search in production


# Step 5: Set up the AutoGen agent with the retrieval tool
async def main() -> None:
    model_client = OpenAIChatCompletionClient(
        model="gpt-4.1",
        api_key=os.environ["OPENAI_API_KEY"],
    )
    agent = AssistantAgent(
        name="contract_assistant",
        model_client=model_client,
        tools=[retrieve_clauses],
        system_message=(
            "You answer questions about a legal contract. "
            "Use retrieve_clauses to find relevant text before answering. "
            "Cite the section name for every statement you make. "
            "Do not invent clause terms that are not present in retrieved passages."
        ),
        # Without this the agent returns the raw tool output instead of an answer
        reflect_on_tool_use=True,
    )

    # Step 6: Run the agent
    await Console(
        agent.run_stream(
            task="What are the termination conditions in this contract, and which section covers them?"
        )
    )
    await model_client.close()


if __name__ == "__main__":
    asyncio.run(main())