diff --git a/main.py b/main.py new file mode 100644 index 0000000..f3bb010 --- /dev/null +++ b/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()