nodes/generic_llm.py

from pathlib import Path
from typing import Any, Callable

from langchain.chat_models import init_chat_model
from langchain.chat_models.base import BaseChatModel
from pydantic import BaseModel

from talentum_score.nodes.base import BaseLLMNode


class GenericLLMNode(BaseLLMNode):
    name: str = "GenericLLMNode"
    description: str = "Generic LLM node"

    def __init__(
        self,
        system_prompt_path: Path | str,
        query: str,
        output_key: str,
        model: str = "gpt-4o-latest",
        system_prompt_kwargs: dict[str, str] = {},
        structured_output: BaseModel | None = None,
        parse_output_fn: Callable[[Any], Any] = lambda x: x["parsed"]
        if isinstance(x, dict) and x.get("parsed") is not None
        else x,
        api_key: str | None = None,
    ):
        self.api_key = api_key
        super().__init__(
            system_prompt_path=system_prompt_path,
            query=query,
            output_key=output_key,
            model=model,
            system_prompt_kwargs=system_prompt_kwargs,
            structured_output=structured_output,
            parse_output_fn=parse_output_fn,
        )

    def get_llm(self, **kwargs) -> BaseChatModel:
        return init_chat_model(
            model=self.model,
            api_key=self.api_key,
            **kwargs,
        )