Package zep_python

class APIError(ZepClientError):
    """
    Raised when the API response format is unexpected.

    Inherits from ZepClientError.
    """

    def __init__(
        self, response: Union[httpx.Response, None] = None, message: str = "API error"
    ) -> None:
        if response:
            response_data = {
                "status_code": response.status_code,
                "message": response.text,
            }
        else:
            response_data = None
        super().__init__(message=message, response_data=response_data)

class Memory(BaseModel):
    """
    Represents a memory object with messages, metadata, and other attributes.

    Attributes
    ----------
    messages : Optional[List[Dict[str, Any]]]
        A list of message objects, where each message contains a role and content.
    metadata : Optional[Dict[str, Any]]
        A dictionary containing metadata associated with the memory.
    summary : Optional[Summary]
        A Summary object.
    uuid : Optional[str]
        A unique identifier for the memory.
    created_at : Optional[str]
        The timestamp when the memory was created.
    token_count : Optional[int]
        The token count of the memory.

    Methods
    -------
    to_dict() -> Dict[str, Any]:
        Returns a dictionary representation of the message.
    """

    messages: List[Message] = Field(
        default=[], description="A List of Messages or empty List is required"
    )
    metadata: Optional[Dict[str, Any]] = Field(optional=True, default=None)
    summary: Optional[Summary] = Field(optional=True, default=None)
    uuid: Optional[str] = Field(optional=True, default=None)
    created_at: Optional[str] = Field(optional=True, default=None)
    token_count: Optional[int] = Field(optional=True, default=None)

    def to_dict(self) -> Dict[str, Any]:
        return self.dict()

class MemorySearchPayload(BaseModel):
    """
    Represents a search payload for querying memory.

    Attributes
    ----------
    metadata : Dict[str, Any]
        Metadata associated with the search query.
    text : str
        The text of the search query.
    search_scope : Optional[str]
        Search over messages or summaries. Defaults to "messages".
        Must be one of "messages" or "summary".
    search_type : Optional[str]
        The type of search to perform. Defaults to "similarity".
        Must be one of "similarity" or "mmr".
    mmr_lambda : Optional[float]
        The lambda parameter for the MMR Reranking Algorithm.
    """

    text: Optional[str] = Field(default=None)
    metadata: Optional[Dict[str, Any]] = Field(default=None)
    search_scope: Optional[str] = Field(default="messages")
    search_type: Optional[str] = Field(default="similarity")
    mmr_lambda: Optional[float] = Field(default=None)

class MemorySearchResult(BaseModel):
    """
    Represents a search result from querying memory.

    Attributes
    ----------
    message : Optional[Dict[str, Any]]
        The message matched by search.
    summary : Optional[Summary]
        The summary matched by search.
    metadata : Optional[Dict[str, Any]]
        Metadata associated with the search result.
    dist : Optional[float]
        The distance metric of the search result.
    """

    # TODO: Legacy bug. message should be a Message object.
    message: Optional[Dict[str, Any]] = None
    summary: Optional[Summary] = None
    metadata: Optional[Dict[str, Any]] = None
    dist: Optional[float] = None

class Message(BaseModel):
    """
    Represents a message in a conversation.

    Attributes
    ----------
    uuid : str, optional
        The unique identifier of the message.
    created_at : str, optional
        The timestamp of when the message was created.
    role : str
        The role of the sender of the message (e.g., "user", "assistant").
    content : str
        The content of the message.
    token_count : int, optional
        The number of tokens in the message.

    Methods
    -------
    to_dict() -> Dict[str, Any]:
        Returns a dictionary representation of the message.
    """

    role: str = Field("A role is required")
    content: str = Field("Content is required")
    uuid: Optional[str] = Field(optional=True, default=None)
    created_at: Optional[str] = Field(optional=True, default=None)
    token_count: Optional[int] = Field(optional=True, default=None)
    metadata: Optional[Dict[str, Any]] = Field(optional=True, default=None)

    def to_dict(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the message.

        Returns
        -------
        Dict[str, Any]
            A dictionary containing the attributes of the message.
        """
        return self.dict()

class NotFoundError(ZepClientError):
    """
    Raised when the API response contains no results.

    Inherits from ZepClientError.
    """

    def __init__(self, message: str) -> None:
        super().__init__(message)

class Session(BaseModel):
    """
    Represents a session object with a unique identifier, metadata,
    and other attributes.

    Attributes
    ----------
    uuid : Optional[str]
        A unique identifier for the session.
        This is generated server-side and is not expected to be present on creation.
    created_at : str
        The timestamp when the session was created.
        Generated by the server.
    updated_at : str
        The timestamp when the session was last updated.
        Generated by the server.
    deleted_at : Optional[datetime]
        The timestamp when the session was deleted.
        Generated by the server.
    session_id : str
        The unique identifier of the session.
    metadata : Dict[str, Any]
        The metadata associated with the session.
    """

    uuid: Optional[str] = None
    id: Optional[int] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None
    deleted_at: Optional[str] = None
    session_id: str
    user_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None

class Summary(BaseModel):
    """
    Represents a summary of a conversation.

    Attributes
    ----------
    uuid : str
        The unique identifier of the summary.
    created_at : str
        The timestamp of when the summary was created.
    content : str
        The content of the summary.
    recent_message_uuid : str
        The unique identifier of the most recent message in the conversation.
    token_count : int
        The number of tokens in the summary.

    Methods
    -------
    to_dict() -> Dict[str, Any]:
        Returns a dictionary representation of the summary.
    """

    uuid: str = Field("A uuid is required")
    created_at: str = Field("A created_at is required")
    content: str = Field("Content is required")
    recent_message_uuid: str = Field("A recent_message_uuid is required")
    token_count: int = Field("A token_count is required")

    def to_dict(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the summary.

        Returns
        -------
        Dict[str, Any]
            A dictionary containing the attributes of the summary.
        """
        return self.dict()

class ZepClient:
    """
    ZepClient class implementation.

    Attributes
    ----------
    base_url : str
        The base URL of the API.
    memory : MemoryClient
        The client used for making Memory API requests.
    document : DocumentClient
        The client used for making Document API requests.

    Methods
    -------
    get_memory(session_id: str, lastn: Optional[int] = None) -> List[Memory]:
        Retrieve memory for the specified session. (Deprecated)
    add_memory(session_id: str, memory_messages: Memory) -> str:
        Add memory to the specified session. (Deprecated)
    delete_memory(session_id: str) -> str:
        Delete memory for the specified session. (Deprecated)
    search_memory(session_id: str, search_payload: SearchPayload,
                  limit: Optional[int] = None) -> List[SearchResult]:
        Search memory for the specified session. (Deprecated)
    close() -> None:
        Close the HTTP client.
    """

    base_url: str
    memory: MemoryClient
    document: DocumentClient
    user: UserClient

    def __init__(self, base_url: str, api_key: Optional[str] = None) -> None:
        """
        Initialize the ZepClient with the specified base URL.

        Parameters
        ----------
        base_url : str
            The base URL of the API.

        api_key : Optional[str]
            The API key to use for authentication. (optional)
        """

        headers: Dict[str, str] = {}
        if api_key is not None:
            headers["Authorization"] = f"Bearer {api_key}"

        self.base_url = concat_url(base_url, API_BASE_PATH)
        self.aclient = httpx.AsyncClient(
            base_url=self.base_url, headers=headers, timeout=API_TIMEOUT
        )
        self.client = httpx.Client(
            base_url=self.base_url, headers=headers, timeout=API_TIMEOUT
        )

        self._healthcheck(base_url)

        self.memory = MemoryClient(self.aclient, self.client)
        self.document = DocumentClient(self.aclient, self.client)
        self.user = UserClient(self.aclient, self.client)

    def _healthcheck(self, base_url: str) -> None:
        """
        Check that the Zep server is running, the API URL is correct,
        and that the server version is compatible with this client.

        Raises
        ------
        ConnectionError
            If the server is not running or the API URL is incorrect.
        """

        url = concat_url(base_url, "/healthz")

        error_msg = """Failed to connect to Zep server. Please check that:
         - the server is running 
         - the API URL is correct
         - No other process is using the same port
         """

        try:
            response = httpx.get(url)
            if response.status_code != 200 or response.text != ".":
                raise APIError(response, error_msg)

            zep_server_version_str = response.headers.get("X-Zep-Version")
            if zep_server_version_str:
                if "dev" in zep_server_version_str:
                    return

                zep_server_version = parse_version_string(zep_server_version_str)
            else:
                zep_server_version = Version("0.0.0")

            if zep_server_version < Version(MINIMUM_SERVER_VERSION):
                warnings.warn(
                    (
                        "You are using an incompatible Zep server version. Please"
                        f" upgrade to {MINIMUM_SERVER_VERSION} or later."
                    ),
                    Warning,
                    stacklevel=2,
                )
        except (httpx.ConnectError, httpx.NetworkError, httpx.TimeoutException) as e:
            raise APIError(None, error_msg) from e

    async def __aenter__(self) -> "ZepClient":
        """Asynchronous context manager entry point"""
        return self

    async def __aexit__(
        self,
        exc_type: Type[Exception],
        exc_val: Exception,
        exc_tb: TracebackType,
    ) -> None:
        """Asynchronous context manager exit point"""
        await self.aclose()

    def __enter__(self) -> "ZepClient":
        """Sync context manager entry point"""
        return self

    def __exit__(
        self,
        exc_type: Type[Exception],
        exc_val: Exception,
        exc_tb: TracebackType,
    ) -> None:
        """Sync context manager exit point"""
        self.close()

    # Facade methods for Memory API
    def get_session(self, session_id: str) -> Session:
        deprecated_warning(self.get_session)
        return self.memory.get_session(session_id)

    async def aget_session(self, session_id: str) -> Session:
        deprecated_warning(self.aget_session)
        return await self.memory.aget_session(session_id)

    def add_session(self, session: Session) -> Session:
        deprecated_warning(self.add_session)
        return self.memory.add_session(session)

    async def aadd_session(self, session: Session) -> Session:
        deprecated_warning(self.aadd_session)
        return await self.memory.aadd_session(session)

    def get_memory(self, session_id: str, lastn: Optional[int] = None) -> Memory:
        deprecated_warning(self.get_memory)
        return self.memory.get_memory(session_id, lastn)

    async def aget_memory(self, session_id: str, lastn: Optional[int] = None) -> Memory:
        deprecated_warning(self.aget_memory)
        return await self.memory.aget_memory(session_id, lastn)

    def add_memory(self, session_id: str, memory_messages: Memory) -> str:
        deprecated_warning(self.add_memory)
        return self.memory.add_memory(session_id, memory_messages)

    async def aadd_memory(self, session_id: str, memory_messages: Memory) -> str:
        deprecated_warning(self.aadd_memory)
        return await self.memory.aadd_memory(session_id, memory_messages)

    def delete_memory(self, session_id: str) -> str:
        deprecated_warning(self.delete_memory)
        return self.memory.delete_memory(session_id)

    async def adelete_memory(self, session_id: str) -> str:
        deprecated_warning(self.adelete_memory)
        return await self.memory.adelete_memory(session_id)

    def search_memory(
        self,
        session_id: str,
        search_payload: MemorySearchPayload,
        limit: Optional[int] = None,
    ) -> List[MemorySearchResult]:
        deprecated_warning(self.search_memory)
        return self.memory.search_memory(session_id, search_payload, limit)

    async def asearch_memory(
        self,
        session_id: str,
        search_payload: MemorySearchPayload,
        limit: Optional[int] = None,
    ) -> List[MemorySearchResult]:
        deprecated_warning(self.asearch_memory)
        return await self.memory.asearch_memory(session_id, search_payload, limit)

    # Close the HTTP client
    async def aclose(self) -> None:
        """
        Asynchronously close the HTTP client.

        [Optional] This method may be called when the ZepClient is no longer needed to
        release resources.
        """
        await self.aclient.aclose()

    def close(self) -> None:
        """
        Close the HTTP client.

        [Optional] This method may be called when the ZepClient is no longer needed to
        release resources.
        """
        self.client.close()