Repository: emcie-co/parlant
File: src/parlant/adapters/nlp/litellm_service.py (Python)

Code before:
            )
            raise


class LiteLLM_Default(LiteLLMSchematicGenerator[T]):
    def __init__(self, logger: Logger) -> None:
        super().__init__(
            model_name=os.environ["LITELLM_PROVIDER_MODEL_NAME"],
            logger=logger,
        )

    @property
    @override
    def max_tokens(self) -> int:
        return 5000

    # 8192 16381


class LiteLLMService(NLPService):
    
    @staticmethod
    def verify_environment() -> str | None:
        """Returns an error message if the environment is not set up correctly."""

        if not os.environ["LITELLM_PROVIDER_MODEL_NAME"]:
            return """\
You're using the LITELLM NLP service, but LITELLM_PROVIDER_MODEL_NAME is not set.
Please set LITELLM_PROVIDER_MODEL_NAME in your environment before running Parlant.
"""
        if not os.environ["LITELLM_PROVIDER_API_KEY"]:
            return """\
You're using the LITELLM NLP service, but LITELLM_PROVIDER_API_KEY is not set.
Please set LITELLM_PROVIDER_API_KEY in your environment before running Parlant.
"""

        return None
    def __init__(
        self,
        logger: Logger,
    ) -> None:
        self._model_name = os.environ["LITELLM_PROVIDER_MODEL_NAME"]
        self._logger = logger
        self._logger.info("Initialized LiteLLMService")

    @override
    async def get_schematic_generator(self, t: type[T]) -> LiteLLMSchematicGenerator[T]:
        return LiteLLM_Default[t](self._logger)  # type: ignore

    @override
    async def get_embedder(self) -> Embedder:

Review comment (bug):
Good work man, just please fix a KeyError bug here, since if there is no `"LITELLM_PROVIDER_MODEL_NAME" in os.environ` than it will raise error.

write either (I prefer the first) `KEY not in os.environ` or `os.environ.get(KEY)` (do to the other env vars obviously)

Code after:
            raise


class LiteLLM_Default(LiteLLMSchematicGenerator[T]):
    def __init__(self, logger: Logger) -> None:
        super().__init__(
            model_name=os.environ.get("LITELLM_PROVIDER_MODEL_NAME"),# type: ignore
            logger=logger,
        )

    @property
    @override
    def max_tokens(self) -> int:
        return 5000

    # 8192 16381


class LiteLLMService(NLPService):
    
    @staticmethod
    def verify_environment() -> str | None:
        """Returns an error message if the environment is not set up correctly."""

        if not os.environ.get("LITELLM_PROVIDER_MODEL_NAME"):
            return """\
You're using the LITELLM NLP service, but LITELLM_PROVIDER_MODEL_NAME is not set.
Please set LITELLM_PROVIDER_MODEL_NAME in your environment before running Parlant.
"""
        if not os.environ.get("LITELLM_PROVIDER_API_KEY"):
            return """\
You're using the LITELLM NLP service, but LITELLM_PROVIDER_API_KEY is not set.
Please set LITELLM_PROVIDER_API_KEY in your environment before running Parlant.
"""

        return None
    
    def __init__(
        self,
        logger: Logger,
    ) -> None:
        self._model_name = os.environ.get("LITELLM_PROVIDER_MODEL_NAME")
        self._logger = logger
        self._logger.info("Initialized LiteLLMService")

    @override
    async def get_schematic_generator(self, t: type[T]) -> LiteLLMSchematicGenerator[T]:
        return LiteLLM_Default[t](self._logger)  # type: ignore

    @override
    async def get_embedder(self) -> Embedder: