graphiti-core 0.4.2__tar.gz → 0.5.0__tar.gz

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.4.2
+Version: 0.5.0
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/cross_encoder/bge_reranker_client.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 import asyncio
-from typing import List, Tuple
 
 from sentence_transformers import CrossEncoder
 
@@ -26,7 +25,7 @@ class BGERerankerClient(CrossEncoderClient):
     def __init__(self):
         self.model = CrossEncoder('BAAI/bge-reranker-v2-m3')
 
-    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
         if not passages:
             return []
 

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/cross_encoder/client.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 from abc import ABC, abstractmethod
-from typing import List, Tuple
 
 
 class CrossEncoderClient(ABC):
@@ -26,16 +25,16 @@ class CrossEncoderClient(ABC):
     """
 
     @abstractmethod
-    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
         """
         Rank the given passages based on their relevance to the query.
 
         Args:
             query (str): The query string.
-            passages (List[str]): A list of passages to rank.
+            passages (list[str]): A list of passages to rank.
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing the passage and its score,
+            list[tuple[str, float]]: A list of tuples containing the passage and its score,
                                      sorted in descending order of relevance.
         """
         pass

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/cross_encoder/openai_reranker_client.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from typing import Any
 
@@ -22,6 +21,7 @@ import openai
 from openai import AsyncOpenAI
 from pydantic import BaseModel
 
+from ..helpers import semaphore_gather
 from ..llm_client import LLMConfig, RateLimitError
 from ..prompts import Message
 from .client import CrossEncoderClient
@@ -75,7 +75,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
             for passage in passages
         ]
         try:
-            responses = await asyncio.gather(
+            responses = await semaphore_gather(
                 *[
                     self.client.chat.completions.create(
                         model=DEFAULT_MODEL,

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/edges.py

@@ -23,6 +23,7 @@ from uuid import uuid4
 
 from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
+from typing_extensions import LiteralString
 
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
@@ -50,7 +51,7 @@ class Edge(BaseModel, ABC):
     async def delete(self, driver: AsyncDriver):
         result = await driver.execute_query(
             """
-        MATCH (n)-[e {uuid: $uuid}]->(m)
+        MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
         DELETE e
         """,
             uuid=self.uuid,
@@ -137,19 +138,35 @@ class EpisodicEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid As uuid,
             e.group_id AS group_id,
             n.uuid AS source_node_uuid, 
             m.uuid AS target_node_uuid, 
             e.created_at AS created_at
-        """,
+        ORDER BY e.uuid DESC 
+        """
+            + limit_query,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -274,11 +291,23 @@ class EntityEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid AS uuid,
             n.uuid AS source_node_uuid,
@@ -292,8 +321,12 @@ class EntityEdge(Edge):
             e.expired_at AS expired_at,
             e.valid_at AS valid_at,
             e.invalid_at AS invalid_at
-        """,
+        ORDER BY e.uuid DESC 
+        """
+            + limit_query,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -365,19 +398,35 @@ class CommunityEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid As uuid,
             e.group_id AS group_id,
             n.uuid AS source_node_uuid, 
             m.uuid AS target_node_uuid, 
             e.created_at AS created_at
-        """,
+        ORDER BY e.uuid DESC
+        """
+            + limit_query,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/embedder/client.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 from abc import ABC, abstractmethod
-from typing import Iterable, List, Literal
+from collections.abc import Iterable
 
 from pydantic import BaseModel, Field
 
@@ -23,12 +23,12 @@ EMBEDDING_DIM = 1024
 
 
 class EmbedderConfig(BaseModel):
-    embedding_dim: Literal[1024] = Field(default=EMBEDDING_DIM, frozen=True)
+    embedding_dim: int = Field(default=EMBEDDING_DIM, frozen=True)
 
 
 class EmbedderClient(ABC):
     @abstractmethod
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         pass

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/embedder/openai.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Iterable, List
+from collections.abc import Iterable
 
 from openai import AsyncOpenAI
 from openai.types import EmbeddingModel
@@ -42,7 +42,7 @@ class OpenAIEmbedder(EmbedderClient):
         self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         result = await self.client.embeddings.create(
             input=input_data, model=self.config.embedding_model

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/embedder/voyage.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Iterable, List
+from collections.abc import Iterable
 
 import voyageai  # type: ignore
 from pydantic import Field
@@ -41,11 +41,11 @@ class VoyageAIEmbedder(EmbedderClient):
         self.client = voyageai.AsyncClient(api_key=config.api_key)
 
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         if isinstance(input_data, str):
             input_list = [input_data]
-        elif isinstance(input_data, List):
+        elif isinstance(input_data, list):
             input_list = [str(i) for i in input_data if i]
         else:
             input_list = [str(i) for i in input_data if i is not None]

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/graphiti.py

@@ -14,9 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
-from datetime import datetime, timezone
+from datetime import datetime
 from time import time
 
 from dotenv import load_dotenv
@@ -27,7 +26,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient, OpenAIClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
@@ -43,10 +42,6 @@ from graphiti_core.search.search_utils import (
     get_relevant_edges,
     get_relevant_nodes,
 )
-from graphiti_core.utils import (
-    build_episodic_edges,
-    retrieve_episodes,
-)
 from graphiti_core.utils.bulk_utils import (
     RawEpisode,
     add_nodes_and_edges_bulk,
@@ -57,12 +52,14 @@ from graphiti_core.utils.bulk_utils import (
     resolve_edge_pointers,
     retrieve_previous_episodes_bulk,
 )
+from graphiti_core.utils.datetime_utils import utc_now
 from graphiti_core.utils.maintenance.community_operations import (
     build_communities,
     remove_communities,
     update_community,
 )
 from graphiti_core.utils.maintenance.edge_operations import (
+    build_episodic_edges,
     dedupe_extracted_edge,
     extract_edges,
     resolve_edge_contradictions,
@@ -71,6 +68,7 @@ from graphiti_core.utils.maintenance.edge_operations import (
 from graphiti_core.utils.maintenance.graph_data_operations import (
     EPISODE_WINDOW_LEN,
     build_indices_and_constraints,
+    retrieve_episodes,
 )
 from graphiti_core.utils.maintenance.node_operations import (
     extract_nodes,
@@ -313,22 +311,26 @@ class Graphiti:
             start = time()
 
             entity_edges: list[EntityEdge] = []
-            now = datetime.now(timezone.utc)
+            now = utc_now()
 
             previous_episodes = await self.retrieve_episodes(
                 reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
             )
-            episode = EpisodicNode(
-                name=name,
-                group_id=group_id,
-                labels=[],
-                source=source,
-                content=episode_body,
-                source_description=source_description,
-                created_at=now,
-                valid_at=reference_time,
+
+            episode = (
+                await EpisodicNode.get_by_uuid(self.driver, uuid)
+                if uuid is not None
+                else EpisodicNode(
+                    name=name,
+                    group_id=group_id,
+                    labels=[],
+                    source=source,
+                    content=episode_body,
+                    source_description=source_description,
+                    created_at=now,
+                    valid_at=reference_time,
+                )
             )
-            episode.uuid = uuid if uuid is not None else episode.uuid
 
             # Extract entities as nodes
 
@@ -337,13 +339,13 @@ class Graphiti:
 
             # Calculate Embeddings
 
-            await asyncio.gather(
+            await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes]
             )
 
             # Find relevant nodes already in the graph
             existing_nodes_lists: list[list[EntityNode]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[get_relevant_nodes(self.driver, [node]) for node in extracted_nodes]
                 )
             )
@@ -351,7 +353,7 @@ class Graphiti:
             # Resolve extracted nodes with nodes already in the graph and extract facts
             logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
-            (mentioned_nodes, uuid_map), extracted_edges = await asyncio.gather(
+            (mentioned_nodes, uuid_map), extracted_edges = await semaphore_gather(
                 resolve_extracted_nodes(
                     self.llm_client,
                     extracted_nodes,
@@ -371,7 +373,7 @@ class Graphiti:
             )
 
             # calculate embeddings
-            await asyncio.gather(
+            await semaphore_gather(
                 *[
                     edge.generate_embedding(self.embedder)
                     for edge in extracted_edges_with_resolved_pointers
@@ -380,7 +382,7 @@ class Graphiti:
 
             # Resolve extracted edges with related edges already in the graph
             related_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -401,7 +403,7 @@ class Graphiti:
             )
 
             existing_source_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -416,7 +418,7 @@ class Graphiti:
             )
 
             existing_target_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -465,7 +467,7 @@ class Graphiti:
 
             # Update any communities
             if update_communities:
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         update_community(self.driver, self.llm_client, self.embedder, node)
                         for node in nodes
@@ -518,7 +520,7 @@ class Graphiti:
         """
         try:
             start = time()
-            now = datetime.now(timezone.utc)
+            now = utc_now()
 
             episodes = [
                 EpisodicNode(
@@ -535,7 +537,7 @@ class Graphiti:
             ]
 
             # Save all the episodes
-            await asyncio.gather(*[episode.save(self.driver) for episode in episodes])
+            await semaphore_gather(*[episode.save(self.driver) for episode in episodes])
 
             # Get previous episode context for each episode
             episode_pairs = await retrieve_previous_episodes_bulk(self.driver, episodes)
@@ -548,19 +550,19 @@ class Graphiti:
             ) = await extract_nodes_and_edges_bulk(self.llm_client, episode_pairs)
 
             # Generate embeddings
-            await asyncio.gather(
+            await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
                 *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
             )
 
             # Dedupe extracted nodes, compress extracted edges
-            (nodes, uuid_map), extracted_edges_timestamped = await asyncio.gather(
+            (nodes, uuid_map), extracted_edges_timestamped = await semaphore_gather(
                 dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes),
                 extract_edge_dates_bulk(self.llm_client, extracted_edges, episode_pairs),
             )
 
             # save nodes to KG
-            await asyncio.gather(*[node.save(self.driver) for node in nodes])
+            await semaphore_gather(*[node.save(self.driver) for node in nodes])
 
             # re-map edge pointers so that they don't point to discard dupe nodes
             extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
@@ -571,7 +573,7 @@ class Graphiti:
             )
 
             # save episodic edges to KG
-            await asyncio.gather(
+            await semaphore_gather(
                 *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers]
             )
 
@@ -584,7 +586,7 @@ class Graphiti:
             # invalidate edges
 
             # save edges to KG
-            await asyncio.gather(*[edge.save(self.driver) for edge in edges])
+            await semaphore_gather(*[edge.save(self.driver) for edge in edges])
 
             end = time()
             logger.info(f'Completed add_episode_bulk in {(end - start) * 1000} ms')
@@ -607,12 +609,12 @@ class Graphiti:
             self.driver, self.llm_client, group_ids
         )
 
-        await asyncio.gather(
+        await semaphore_gather(
             *[node.generate_name_embedding(self.embedder) for node in community_nodes]
         )
 
-        await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
-        await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
+        await semaphore_gather(*[node.save(self.driver) for node in community_nodes])
+        await semaphore_gather(*[edge.save(self.driver) for edge in community_edges])
 
         return community_nodes
 
@@ -695,7 +697,7 @@ class Graphiti:
     async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
         episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
-        edges_list = await asyncio.gather(
+        edges_list = await semaphore_gather(
             *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
         )
 

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/helpers.py

@@ -14,7 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
 import os
+from collections.abc import Coroutine
 from datetime import datetime
 
 import numpy as np
@@ -25,7 +27,9 @@ load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
+SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
 MAX_REFLEXION_ITERATIONS = 2
+DEFAULT_PAGE_LIMIT = 20
 
 
 def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
@@ -56,6 +60,12 @@ def lucene_sanitize(query: str) -> str:
             ':': r'\:',
             '\\': r'\\',
             '/': r'\/',
+            'O': r'\O',
+            'R': r'\R',
+            'N': r'\N',
+            'T': r'\T',
+            'A': r'\A',
+            'D': r'\D',
         }
     )
 
@@ -73,3 +83,19 @@ def normalize_l2(embedding: list[float]) -> list[float]:
     else:
         norm = np.linalg.norm(embedding_array, 2, axis=1, keepdims=True)
         return (np.where(norm == 0, embedding_array, embedding_array / norm)).tolist()
+
+
+# Use this instead of asyncio.gather() to bound coroutines
+async def semaphore_gather(
+    *coroutines: Coroutine, max_coroutines: int = SEMAPHORE_LIMIT, return_exceptions=True
+):
+    semaphore = asyncio.Semaphore(max_coroutines)
+
+    async def _wrap_coroutine(coroutine):
+        async with semaphore:
+            return await coroutine
+
+    return await asyncio.gather(
+        *(_wrap_coroutine(coroutine) for coroutine in coroutines),
+        return_exceptions=return_exceptions,
+    )

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/llm_client/anthropic_client.py

@@ -20,6 +20,7 @@ import typing
 
 import anthropic
 from anthropic import AsyncAnthropic
+from pydantic import BaseModel
 
 from ..prompts.models import Message
 from .client import LLMClient
@@ -46,7 +47,9 @@ class AnthropicClient(LLMClient):
             max_retries=1,
         )
 
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         system_message = messages[0]
         user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]] + [
             {'role': 'assistant', 'content': '{'}

{graphiti_core-0.4.2 → graphiti_core-0.5.0}/graphiti_core/llm_client/client.py

@@ -22,6 +22,7 @@ from abc import ABC, abstractmethod
 
 import httpx
 from diskcache import Cache
+from pydantic import BaseModel
 from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
 
 from ..prompts.models import Message
@@ -55,6 +56,28 @@ class LLMClient(ABC):
         self.cache_enabled = cache
         self.cache_dir = Cache(DEFAULT_CACHE_DIR)  # Create a cache directory
 
+    def _clean_input(self, input: str) -> str:
+        """Clean input string of invalid unicode and control characters.
+
+        Args:
+            input: Raw input string to be cleaned
+
+        Returns:
+            Cleaned string safe for LLM processing
+        """
+        # Clean any invalid Unicode
+        cleaned = input.encode('utf-8', errors='ignore').decode('utf-8')
+
+        # Remove zero-width characters and other invisible unicode
+        zero_width = '\u200b\u200c\u200d\ufeff\u2060'
+        for char in zero_width:
+            cleaned = cleaned.replace(char, '')
+
+        # Remove control characters except newlines, returns, and tabs
+        cleaned = ''.join(char for char in cleaned if ord(char) >= 32 or char in '\n\r\t')
+
+        return cleaned
+
     @retry(
         stop=stop_after_attempt(4),
         wait=wait_random_exponential(multiplier=10, min=5, max=120),
@@ -66,14 +89,18 @@ class LLMClient(ABC):
         else None,
         reraise=True,
     )
-    async def _generate_response_with_retry(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response_with_retry(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         try:
-            return await self._generate_response(messages)
+            return await self._generate_response(messages, response_model)
         except (httpx.HTTPStatusError, RateLimitError) as e:
             raise e
 
     @abstractmethod
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         pass
 
     def _get_cache_key(self, messages: list[Message]) -> str:
@@ -82,7 +109,17 @@ class LLMClient(ABC):
         key_str = f'{self.model}:{message_str}'
         return hashlib.md5(key_str.encode()).hexdigest()
 
-    async def generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
+        if response_model is not None:
+            serialized_model = json.dumps(response_model.model_json_schema())
+            messages[
+                -1
+            ].content += (
+                f'\n\nRespond with a JSON object in the following format:\n\n{serialized_model}'
+            )
+
         if self.cache_enabled:
             cache_key = self._get_cache_key(messages)
 
@@ -91,7 +128,10 @@ class LLMClient(ABC):
                 logger.debug(f'Cache hit for {cache_key}')
                 return cached_response
 
-        response = await self._generate_response_with_retry(messages)
+        for message in messages:
+            message.content = self._clean_input(message.content)
+
+        response = await self._generate_response_with_retry(messages, response_model)
 
         if self.cache_enabled:
             self.cache_dir.set(cache_key, response)