Python Modules (Developer)

APIAgentConfiguration

Configuration for the Lex Machina API Agent. It supports authentication via API token, OAuth2 client credentials, or delegation URL.

Source code in src/lexmachina_agent/agent_executor.py

class APIAgentConfiguration:
    """Configuration for the Lex Machina API Agent.
    It supports authentication via API token, OAuth2 client credentials, or delegation URL."""

    def __init__(self) -> None:
        # Load configuration from environment variables
        api_base_url = os.environ.get("API_BASE_URL", "https://law-api-poc.stage.lexmachina.com")
        token = os.environ.get("API_TOKEN")
        client_id = os.environ.get("CLIENT_ID")
        client_secret = os.environ.get("CLIENT_SECRET")
        delegation_url = os.environ.get("DELEGATION_URL")

        if all(v is None for v in [token, client_id, client_secret, delegation_url]):
            raise MissingConfigurationError(["API_TOKEN", "CLIENT_ID", "CLIENT_SECRET", "DELEGATION_URL"])

        if token:
            logger.warning(
                "Using API_TOKEN for authentication. Consider using CLIENT_ID / CLIENT_SECRET, or DELEGATION_URL for better security."
            )
        if client_id is not None and client_secret is None:
            raise RequiredConfigurationError("CLIENT_SECRET")
        if client_secret is not None and client_id is None:
            raise RequiredConfigurationError("CLIENT_ID")

        self.api_base_url = api_base_url
        self.token = token
        self.client_id = client_id
        self.client_secret = client_secret
        self.delegation_url = delegation_url

    @property
    def is_using_delegation(self) -> bool:
        return self.delegation_url is not None

    def build_agent(self) -> "LexMachinaAPIAgent":
        """Constructs and returns a LexMachinaAPIAgent instance based on the configuration."""
        if self.token:
            return LexMachinaAPIAgent(
                api_base_url=self.api_base_url,
                token=self.token,
            )
        elif self.client_id and self.client_secret:
            data = {
                "grant_type": "client_credentials",
                "client_id": self.client_id,
                "client_secret": self.client_secret,
            }
            token_url = f"{self.api_base_url}/api/token"
            try:
                resp = httpx.post(token_url, data=data, headers={"Accept": "application/json"})
                resp.raise_for_status()
                access_token = resp.json().get("access_token")
                if not access_token:
                    logger.error("Token endpoint did not return access_token.")
                    raise ConfigurationError()
                return LexMachinaAPIAgent(
                    api_base_url=self.api_base_url,
                    token=typing.cast(str, access_token),
                )
            except httpx.HTTPError:
                logger.exception("OAuth2 token request failed.")
                raise
        elif self.delegation_url:
            # Implement delegation URL based authentication
            raise NotImplementedError("Delegation URL authentication not implemented yet.")
        else:
            raise ConfigurationError()  # This should not happen

build_agent()

Constructs and returns a LexMachinaAPIAgent instance based on the configuration.

Source code in src/lexmachina_agent/agent_executor.py

def build_agent(self) -> "LexMachinaAPIAgent":
    """Constructs and returns a LexMachinaAPIAgent instance based on the configuration."""
    if self.token:
        return LexMachinaAPIAgent(
            api_base_url=self.api_base_url,
            token=self.token,
        )
    elif self.client_id and self.client_secret:
        data = {
            "grant_type": "client_credentials",
            "client_id": self.client_id,
            "client_secret": self.client_secret,
        }
        token_url = f"{self.api_base_url}/api/token"
        try:
            resp = httpx.post(token_url, data=data, headers={"Accept": "application/json"})
            resp.raise_for_status()
            access_token = resp.json().get("access_token")
            if not access_token:
                logger.error("Token endpoint did not return access_token.")
                raise ConfigurationError()
            return LexMachinaAPIAgent(
                api_base_url=self.api_base_url,
                token=typing.cast(str, access_token),
            )
        except httpx.HTTPError:
            logger.exception("OAuth2 token request failed.")
            raise
    elif self.delegation_url:
        # Implement delegation URL based authentication
        raise NotImplementedError("Delegation URL authentication not implemented yet.")
    else:
        raise ConfigurationError()  # This should not happen

ConfigurationError

Bases: Exception

Base class for configuration-related errors.

Source code in src/lexmachina_agent/agent_executor.py

class ConfigurationError(Exception):
    """Base class for configuration-related errors."""

    def __init__(self) -> None:
        self.args = ("Invalid configuration",)

LexMachinaAPIAgent

A stateful agent that manages communication with the external Protégé in Lex Machina Agent API. It holds the HTTP client and authentication token.

Source code in src/lexmachina_agent/agent_executor.py

class LexMachinaAPIAgent:
    """
    A stateful agent that manages communication with the external Protégé in Lex Machina Agent API.
    It holds the HTTP client and authentication token.
    """

    def __init__(self, api_base_url: str, token: str):
        self._api_base_url = api_base_url
        self._headers = {
            "Authorization": f"Bearer {token}",
            "Accept": "application/json",
        }
        self._client = httpx.AsyncClient(base_url=self._api_base_url, headers=self._headers)
        logger.info("LexMachinaAPIAgent initialized.")

    async def get_suggested_searches(self, query: str) -> dict:
        """Calls the /search/ai_suggested endpoint."""
        try:
            logger.info(f"Fetching suggested searches for: '{query}'")
            response = await self._client.get("/search/ai_suggested", params={"q": query})
            response.raise_for_status()
            return typing.cast(dict, response.json())
        except httpx.HTTPStatusError as e:
            logger.exception("API Error")
            return {"error": str(e)}

    async def get_search_description(self, description_url: str) -> dict:
        """Fetches the description for a single suggested search."""
        try:
            logger.debug(f"Fetching description from: {description_url}")
            response = await self._client.get(description_url)
            response.raise_for_status()
            return typing.cast(dict, response.json())
        except httpx.HTTPStatusError as e:
            logger.exception("API Error")
            return {"error": str(e)}

    async def process_query(self, query: str) -> dict:
        """
        Main method to process a query, fetch suggestions, and enrich them in parallel.
        """
        # 1. Get initial suggestions from the API agent
        suggestions_response = await self.get_suggested_searches(query)

        if "error" in suggestions_response or not suggestions_response.get("result"):
            logger.error("Failed to get initial suggestions.")
            return {
                "error": "Failed to get initial suggestions.",
                "details": suggestions_response,
            }

        # 2. Prepare for parallel enrichment
        suggestions = suggestions_response["result"]
        enrichment_tasks = []
        for suggestion in suggestions:
            # For each suggestion, create a remote call to fetch its description
            task = self.get_search_description(suggestion["description_url"])
            enrichment_tasks.append(task)

        logger.debug(f"Fetching {len(enrichment_tasks)} descriptions in parallel...")

        # 3. Execute all enrichment tasks concurrently and gather results
        descriptions = await asyncio.gather(*enrichment_tasks)

        # 4. Combine the original suggestions with their fetched descriptions
        for suggestion, description_data in zip(suggestions, descriptions):
            suggestion["enriched_description"] = description_data

        logger.debug("Query processing complete")
        return suggestions_response

get_search_description(description_url) async

Fetches the description for a single suggested search.

Source code in src/lexmachina_agent/agent_executor.py

async def get_search_description(self, description_url: str) -> dict:
    """Fetches the description for a single suggested search."""
    try:
        logger.debug(f"Fetching description from: {description_url}")
        response = await self._client.get(description_url)
        response.raise_for_status()
        return typing.cast(dict, response.json())
    except httpx.HTTPStatusError as e:
        logger.exception("API Error")
        return {"error": str(e)}

get_suggested_searches(query) async

Calls the /search/ai_suggested endpoint.

Source code in src/lexmachina_agent/agent_executor.py

async def get_suggested_searches(self, query: str) -> dict:
    """Calls the /search/ai_suggested endpoint."""
    try:
        logger.info(f"Fetching suggested searches for: '{query}'")
        response = await self._client.get("/search/ai_suggested", params={"q": query})
        response.raise_for_status()
        return typing.cast(dict, response.json())
    except httpx.HTTPStatusError as e:
        logger.exception("API Error")
        return {"error": str(e)}

process_query(query) async

Main method to process a query, fetch suggestions, and enrich them in parallel.

Source code in src/lexmachina_agent/agent_executor.py

async def process_query(self, query: str) -> dict:
    """
    Main method to process a query, fetch suggestions, and enrich them in parallel.
    """
    # 1. Get initial suggestions from the API agent
    suggestions_response = await self.get_suggested_searches(query)

    if "error" in suggestions_response or not suggestions_response.get("result"):
        logger.error("Failed to get initial suggestions.")
        return {
            "error": "Failed to get initial suggestions.",
            "details": suggestions_response,
        }

    # 2. Prepare for parallel enrichment
    suggestions = suggestions_response["result"]
    enrichment_tasks = []
    for suggestion in suggestions:
        # For each suggestion, create a remote call to fetch its description
        task = self.get_search_description(suggestion["description_url"])
        enrichment_tasks.append(task)

    logger.debug(f"Fetching {len(enrichment_tasks)} descriptions in parallel...")

    # 3. Execute all enrichment tasks concurrently and gather results
    descriptions = await asyncio.gather(*enrichment_tasks)

    # 4. Combine the original suggestions with their fetched descriptions
    for suggestion, description_data in zip(suggestions, descriptions):
        suggestion["enriched_description"] = description_data

    logger.debug("Query processing complete")
    return suggestions_response

LexmachinaAgentExecutor

Bases: AgentExecutor

AgentExecutor implementation for the Lex Machina agent.

Source code in src/lexmachina_agent/agent_executor.py

class LexmachinaAgentExecutor(AgentExecutor):
    """AgentExecutor implementation for the Lex Machina agent."""

    def __init__(self, config: APIAgentConfiguration) -> None:
        self.config = config

    async def execute(
        self,
        context: RequestContext,
        event_queue: EventQueue,
    ) -> None:
        api = self.config.build_agent()
        if context.task_id is None or context.context_id is None or context.message is None:
            raise ServerError(error=InvalidParamsError(message="Missing task_id or context_id or message"))
        query = context.get_user_input()

        results = await api.process_query(query)
        parts = [Part(root=TextPart(text=str(results)))]
        await event_queue.enqueue_event(
            completed_task(
                context.task_id,
                context.context_id,
                [new_artifact(parts, f"suggestion_{context.task_id}")],
                [context.message],
            )
        )

    async def cancel(self, request: RequestContext, event_queue: EventQueue) -> None:
        """Cancellation is not supported."""
        raise ServerError(error=UnsupportedOperationError())

cancel(request, event_queue) async

Cancellation is not supported.

Source code in src/lexmachina_agent/agent_executor.py

async def cancel(self, request: RequestContext, event_queue: EventQueue) -> None:
    """Cancellation is not supported."""
    raise ServerError(error=UnsupportedOperationError())

MissingConfigurationError

Bases: ConfigurationError

Raised when all configuration fields are missing.

Source code in src/lexmachina_agent/agent_executor.py

class MissingConfigurationError(ConfigurationError):
    """Raised when all configuration fields are missing."""

    def __init__(self, missing_fields: list[str]) -> None:
        self.args = (f"Missing configuration values: {', '.join(missing_fields)}",)

RequiredConfigurationError

Bases: ConfigurationError

Raised when a required configuration field is missing.

Source code in src/lexmachina_agent/agent_executor.py

class RequiredConfigurationError(ConfigurationError):
    """Raised when a required configuration field is missing."""

    def __init__(self, field_name: str) -> None:
        self.args = (f"Missing required configuration value: {field_name}",)

Server application for the Lex Machina A2A agent proxy

app()

Create the Starlette ASGI application with the Lex Machina agent.

Source code in src/lexmachina_agent/server.py

def app() -> Starlette:
    """Create the Starlette ASGI application with the Lex Machina agent."""
    base_url = os.environ.get("BASE_URL", "http://localhost:10011")
    config = APIAgentConfiguration()
    capabilities = AgentCapabilities(streaming=False)
    skill = AgentSkill(
        id="search_suggestions",
        name="Search Suggestions",
        description="""This takes a natural language question or prompt and suggests options for searches.
          You can ask for assistance in building searches and finding the analytics in Lex Machina you care about.""",
        tags=["search", "suggestions", "analytics", "legal"],
        examples=[
            "What is the average time to resolution for contracts cases in SDNY in the last 3 months?",
            "Time to trial in a Los Angeles County case before Judge Randy Rhodes?",
            "Reversal rate for employment cases in the 5th circuit?",
            "Patent cases that went to trial in the last 90 days",
            "Cases before Judge Schofield that mention tortious interference",
            "Has Warby Parker been sued in Texas?",
            "Complaints in Torts cases that mention section 552",
            "Pleadings that mention jurisprudence",
            "Jury Verdicts filed in California in the last 5 years",
            "How long do LA contracts cases before Judge Katherine Chilton take to get to trial?",
            "What is Judge Lemelle's grant rate for transfer motions?",
            "Which firms have the most experience arguing employment cases in N.D.Ill?",
        ],
    )
    agent_card = AgentCard(
        name="Protégé in Lex Machina Agent",
        description="Provide search suggestions based on user input.",
        url=base_url,
        version="1.0.0",
        default_input_modes=["text"],
        default_output_modes=["application/json"],
        capabilities=capabilities,
        skills=[skill],
    )

    request_handler = DefaultRequestHandler(
        agent_executor=LexmachinaAgentExecutor(config=config),
        task_store=InMemoryTaskStore(),
    )

    server = A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)
    return server.build()