127 lines
5.0 KiB
Python
127 lines
5.0 KiB
Python
"""
|
||
Structured output with Union events API assignment.
|
||
|
||
This script demonstrates how to parse raw log lines into typed events using LangChain's
|
||
structured output capabilities and Pydantic v2. It supports two event types:
|
||
|
||
* ``HttpOkEvent`` – successful HTTP request (status 200)
|
||
* ``HttpErrorEvent`` – error response (4xx/5xx)
|
||
|
||
The script can be used as a CLI tool: provide a log file or paste the log into stdin.
|
||
It prints a table with the parsed events.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sys
|
||
from typing import List, Union
|
||
|
||
from langchain_core.output_parsers import PydanticOutputParser
|
||
from langchain_openai import ChatOpenAI
|
||
from pydantic import BaseModel, Field, Literal
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. Define the event models
|
||
# ---------------------------------------------------------------------------
|
||
class HttpOkEvent(BaseModel):
|
||
kind: Literal["ok"] = Field("ok", description="Discriminator for OK events")
|
||
status: int = Field(..., description="HTTP status code (200)")
|
||
path: str = Field(..., description="Request path")
|
||
duration_ms: int = Field(..., description="Duration in milliseconds")
|
||
|
||
class HttpErrorEvent(BaseModel):
|
||
kind: Literal["error"] = Field("error", description="Discriminator for error events")
|
||
status: int = Field(..., description="HTTP status code (4xx/5xx)")
|
||
path: str = Field(..., description="Request path")
|
||
error_message: str = Field(..., description="Error message from the server")
|
||
|
||
# Union with discriminator
|
||
ApiEvent = Union[HttpOkEvent, HttpErrorEvent]
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Create a parser that will be used by LangChain
|
||
# ---------------------------------------------------------------------------
|
||
parser = PydanticOutputParser(pydantic_object=ApiEvent)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. Build a simple chain: prompt -> LLM -> parser
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini") # adjust as needed
|
||
|
||
prompt_template = (
|
||
"Parse the following log line into a JSON object that matches one of the\n"
|
||
"following schemas:\n"
|
||
f"{HttpOkEvent.model_json_schema()}\n"
|
||
f"{HttpErrorEvent.model_json_schema()}\n"
|
||
"The output must be valid JSON and contain only the fields defined in the schema.\n"
|
||
"Do not add any additional keys or text.")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Helper to parse a single line using the LLM chain
|
||
# ---------------------------------------------------------------------------
|
||
from langchain_core.prompts import PromptTemplate
|
||
from langchain.chains import LLMChain
|
||
|
||
prompt = PromptTemplate.from_template(prompt_template)
|
||
chain = LLMChain(llm=llm, prompt=prompt, output_parser=parser)
|
||
|
||
|
||
def parse_line(line: str) -> ApiEvent:
|
||
"""Parse a single log line into an ``ApiEvent`` using the LLM chain."""
|
||
return chain.invoke({"input": line}) # type: ignore[arg-type]
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. CLI handling
|
||
# ---------------------------------------------------------------------------
|
||
def read_input(file_path: str | None) -> List[str]:
|
||
if file_path:
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
return [l.strip() for l in f.readlines() if l.strip()]
|
||
# Read from stdin
|
||
data = sys.stdin.read()
|
||
return [l.strip() for l in data.splitlines() if l.strip()]
|
||
|
||
|
||
def main(argv: List[str] | None = None) -> None:
|
||
parser_cli = argparse.ArgumentParser(description="Parse raw log into typed events")
|
||
parser_cli.add_argument("-f", "--file", help="Path to a file containing the log. If omitted, read from stdin.")
|
||
args = parser_cli.parse_args(argv)
|
||
|
||
lines = read_input(args.file)
|
||
if not lines:
|
||
print("No input provided.", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
events: List[ApiEvent] = []
|
||
for line in lines:
|
||
try:
|
||
event = parse_line(line)
|
||
events.append(event)
|
||
except Exception as exc: # pragma: no cover - LLM errors are runtime
|
||
print(f"Failed to parse line: {line!r}\nError: {exc}", file=sys.stderr)
|
||
continue
|
||
|
||
# Print table
|
||
header = ["kind", "path", "status", "duration_ms", "error_message"]
|
||
rows = []
|
||
for ev in events:
|
||
if isinstance(ev, HttpOkEvent):
|
||
rows.append([ev.kind, ev.path, ev.status, ev.duration_ms, ""]) # type: ignore[arg-type]
|
||
else:
|
||
rows.append([ev.kind, ev.path, ev.status, "", ev.error_message]) # type: ignore[arg-type]
|
||
|
||
# Simple pretty print
|
||
col_widths = [max(len(str(row[i])) for row in ([header] + rows)) for i in range(len(header))]
|
||
fmt = " | ".join(f"{{:<{w}}}" for w in col_widths)
|
||
sep = "-+-".join("-" * w for w in col_widths)
|
||
|
||
print(fmt.format(*header))
|
||
print(sep)
|
||
for row in rows:
|
||
print(fmt.format(*row))
|
||
|
||
|
||
if __name__ == "__main__": # pragma: no cover - entry point
|
||
main()
|