Add main.py

main.py

@@ -0,0 +1,126 @@
+"""
+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()