Повторный экзамен: Structured output — Union событий API: solution.py

solution.py

@@ -0,0 +1,130 @@
+import os
+import argparse
+from typing import Annotated, Literal, Union, List
+
+from pydantic import BaseModel, Field, ValidationError
+from langchain_core.output_parsers import PydanticOutputParser
+from langchain_core.prompts import PromptTemplate
+from langchain_openai import ChatOpenAI
+
+# --------------------------------------------------------------------------- #
+# 1. Модели событий
+# --------------------------------------------------------------------------- #
+
+class HttpOkEvent(BaseModel):
+    kind: Literal["ok"] = Field(..., description="Тип события: OK")
+    status: Literal[200] = Field(..., description="HTTP статус 200")
+    path: str = Field(..., description="Запрошенный путь")
+    duration_ms: int = Field(..., description="Время выполнения в миллисекундах")
+
+class HttpErrorEvent(BaseModel):
+    kind: Literal["error"] = Field(..., description="Тип события: ошибка")
+    status: int = Field(..., description="HTTP статус 4xx/5xx")
+    path: str = Field(..., description="Запрошенный путь")
+    error_message: str = Field(..., description="Текст ошибки")
+
+ApiEvent = Annotated[
+    Union[HttpOkEvent, HttpErrorEvent],
+    Field(discriminator="kind")
+]
+
+# --------------------------------------------------------------------------- #
+# 2. Парсер LLM + Pydantic
+# --------------------------------------------------------------------------- #
+
+parser = PydanticOutputParser(pydantic_object=ApiEvent)
+
+prompt_template = PromptTemplate(
+    input_variables=["log_line"],
+    template=(
+        "Найди в следующем логе событие API и выведи его в формате JSON, "
+        "соответствующем одной из следующих схем:\n"
+        "1. {ok_schema}\n"
+        "2. {error_schema}\n"
+        "Укажи поле \"kind\" как \"ok\" или \"error\".\n"
+        "Лог: {log_line}\n"
+        "Ответ в чистом JSON без комментариев."
+    ),
+)
+
+# --------------------------------------------------------------------------- #
+# 3. Функция обработки одной строки
+# --------------------------------------------------------------------------- #
+
+def parse_line(line: str, llm: ChatOpenAI) -> ApiEvent:
+    ok_schema = HttpOkEvent.model_json_schema()
+    error_schema = HttpErrorEvent.model_json_schema()
+    prompt = prompt_template.format(
+        log_line=line,
+        ok_schema=ok_schema,
+        error_schema=error_schema,
+    )
+    response = llm.invoke(prompt)
+    try:
+        event = parser.parse(response.content)
+    except ValidationError as e:
+        raise ValueError(f"LLM не смог распарсить строку: {line}\n{e}") from e
+    return event
+
+# --------------------------------------------------------------------------- #
+# 4. Обработка всего лога
+# --------------------------------------------------------------------------- #
+
+def parse_log(text: str, llm: ChatOpenAI) -> List[ApiEvent]:
+    events: List[ApiEvent] = []
+    # Разделяем по строкам, игнорируем пустые
+    for line in text.strip().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        event = parse_line(line, llm)
+        events.append(event)
+    return events
+
+# --------------------------------------------------------------------------- #
+# 5. CLI
+# --------------------------------------------------------------------------- #
+
+def main():
+    parser_cli = argparse.ArgumentParser(description="Parse API logs into structured events.")
+    parser_cli.add_argument(
+        "--log",
+        type=str,
+        default=None,
+        help="Путь к файлу с логом. Если не задан, используется пример.",
+    )
+    args = parser_cli.parse_args()
+
+    if args.log:
+        with open(args.log, "r", encoding="utf-8") as f:
+            log_text = f.read()
+    else:
+        log_text = """\
+2023-10-01 12:00:01 INFO /api/users 200 123ms
+2023-10-01 12:00:02 ERROR /api/orders 404 Not Found
+2023-10-01 12:00:03 WARN /api/payments 500 Internal Server Error
+"""
+
+    # Инициализируем LLM
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    if not openai_api_key:
+        raise RuntimeError("Переменная окружения OPENAI_API_KEY не найдена")
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=openai_api_key)
+
+    events = parse_log(log_text, llm)
+
+    # Выводим каждый объект
+    for ev in events:
+        print(ev.model_dump())
+
+    # Таблица
+    print("\n| kind | path | status |")
+    print("|------|------|--------|")
+    for ev in events:
+        if ev.kind == "ok":
+            print(f"| {ev.kind} | {ev.path} | {ev.status} |")
+        else:
+            print(f"| {ev.kind} | {ev.path} | {ev.status} |")
+
+if __name__ == "__main__":
+    main()