graphiti-core 0.1.0__py3-none-any.whl

graphiti_core/utils/maintenance/edge_operations.py

@@ -0,0 +1,170 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+from datetime import datetime
+from time import time
+from typing import List
+
+from graphiti_core.edges import EntityEdge, EpisodicEdge
+from graphiti_core.llm_client import LLMClient
+from graphiti_core.nodes import EntityNode, EpisodicNode
+from graphiti_core.prompts import prompt_library
+
+logger = logging.getLogger(__name__)
+
+
+def build_episodic_edges(
+    entity_nodes: List[EntityNode],
+    episode: EpisodicNode,
+    created_at: datetime,
+) -> List[EpisodicEdge]:
+    edges: List[EpisodicEdge] = []
+
+    for node in entity_nodes:
+        edge = EpisodicEdge(
+            source_node_uuid=episode.uuid,
+            target_node_uuid=node.uuid,
+            created_at=created_at,
+        )
+        edges.append(edge)
+
+    return edges
+
+
+async def extract_edges(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+    nodes: list[EntityNode],
+    previous_episodes: list[EpisodicNode],
+) -> list[EntityEdge]:
+    start = time()
+
+    # Prepare context for LLM
+    context = {
+        'episode_content': episode.content,
+        'episode_timestamp': (episode.valid_at.isoformat() if episode.valid_at else None),
+        'nodes': [
+            {'uuid': node.uuid, 'name': node.name, 'summary': node.summary} for node in nodes
+        ],
+        'previous_episodes': [
+            {
+                'content': ep.content,
+                'timestamp': ep.valid_at.isoformat() if ep.valid_at else None,
+            }
+            for ep in previous_episodes
+        ],
+    }
+
+    llm_response = await llm_client.generate_response(prompt_library.extract_edges.v2(context))
+    print(llm_response)
+    edges_data = llm_response.get('edges', [])
+
+    end = time()
+    logger.info(f'Extracted new edges: {edges_data} in {(end - start) * 1000} ms')
+
+    # Convert the extracted data into EntityEdge objects
+    edges = []
+    for edge_data in edges_data:
+        if edge_data['target_node_uuid'] and edge_data['source_node_uuid']:
+            edge = EntityEdge(
+                source_node_uuid=edge_data['source_node_uuid'],
+                target_node_uuid=edge_data['target_node_uuid'],
+                name=edge_data['relation_type'],
+                fact=edge_data['fact'],
+                episodes=[episode.uuid],
+                created_at=datetime.now(),
+                valid_at=None,
+                invalid_at=None,
+            )
+            edges.append(edge)
+            logger.info(
+                f'Created new edge: {edge.name} from (UUID: {edge.source_node_uuid}) to (UUID: {edge.target_node_uuid})'
+            )
+
+    return edges
+
+
+def create_edge_identifier(
+    source_node: EntityNode, edge: EntityEdge, target_node: EntityNode
+) -> str:
+    return f'{source_node.name}-{edge.name}-{target_node.name}'
+
+
+async def dedupe_extracted_edges(
+    llm_client: LLMClient,
+    extracted_edges: list[EntityEdge],
+    existing_edges: list[EntityEdge],
+) -> list[EntityEdge]:
+    # Create edge map
+    edge_map = {}
+    for edge in extracted_edges:
+        edge_map[edge.uuid] = edge
+
+    # Prepare context for LLM
+    context = {
+        'extracted_edges': [
+            {'uuid': edge.uuid, 'name': edge.name, 'fact': edge.fact} for edge in extracted_edges
+        ],
+        'existing_edges': [
+            {'uuid': edge.uuid, 'name': edge.name, 'fact': edge.fact} for edge in existing_edges
+        ],
+    }
+
+    llm_response = await llm_client.generate_response(prompt_library.dedupe_edges.v1(context))
+    unique_edge_data = llm_response.get('unique_facts', [])
+    logger.info(f'Extracted unique edges: {unique_edge_data}')
+
+    # Get full edge data
+    edges = []
+    for unique_edge in unique_edge_data:
+        edge = edge_map[unique_edge['uuid']]
+        edges.append(edge)
+
+    return edges
+
+
+async def dedupe_edge_list(
+    llm_client: LLMClient,
+    edges: list[EntityEdge],
+) -> list[EntityEdge]:
+    start = time()
+
+    # Create edge map
+    edge_map = {}
+    for edge in edges:
+        edge_map[edge.uuid] = edge
+
+    # Prepare context for LLM
+    context = {'edges': [{'uuid': edge.uuid, 'fact': edge.fact} for edge in edges]}
+
+    llm_response = await llm_client.generate_response(
+        prompt_library.dedupe_edges.edge_list(context)
+    )
+    unique_edges_data = llm_response.get('unique_facts', [])
+
+    end = time()
+    logger.info(f'Extracted edge duplicates: {unique_edges_data} in {(end - start) * 1000} ms ')
+
+    # Get full edge data
+    unique_edges = []
+    for edge_data in unique_edges_data:
+        uuid = edge_data['uuid']
+        edge = edge_map[uuid]
+        edge.fact = edge_data['fact']
+        unique_edges.append(edge)
+
+    return unique_edges

graphiti_core/utils/maintenance/graph_data_operations.py

@@ -0,0 +1,133 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import asyncio
+import logging
+from datetime import datetime, timezone
+
+from neo4j import AsyncDriver
+from typing_extensions import LiteralString
+
+from graphiti_core.nodes import EpisodeType, EpisodicNode
+
+EPISODE_WINDOW_LEN = 3
+
+logger = logging.getLogger(__name__)
+
+
+async def build_indices_and_constraints(driver: AsyncDriver):
+    range_indices: list[LiteralString] = [
+        'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
+        'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
+        'CREATE INDEX relation_uuid IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.uuid)',
+        'CREATE INDEX mention_uuid IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.uuid)',
+        'CREATE INDEX name_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.name)',
+        'CREATE INDEX created_at_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.created_at)',
+        'CREATE INDEX created_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.created_at)',
+        'CREATE INDEX valid_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.valid_at)',
+        'CREATE INDEX name_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.name)',
+        'CREATE INDEX created_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.created_at)',
+        'CREATE INDEX expired_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.expired_at)',
+        'CREATE INDEX valid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.valid_at)',
+        'CREATE INDEX invalid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.invalid_at)',
+    ]
+
+    fulltext_indices: list[LiteralString] = [
+        'CREATE FULLTEXT INDEX name_and_summary IF NOT EXISTS FOR (n:Entity) ON EACH [n.name, n.summary]',
+        'CREATE FULLTEXT INDEX name_and_fact IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact]',
+    ]
+
+    vector_indices: list[LiteralString] = [
+        """
+        CREATE VECTOR INDEX fact_embedding IF NOT EXISTS
+        FOR ()-[r:RELATES_TO]-() ON (r.fact_embedding)
+        OPTIONS {indexConfig: {
+         `vector.dimensions`: 1024,
+         `vector.similarity_function`: 'cosine'
+        }}
+        """,
+        """
+        CREATE VECTOR INDEX name_embedding IF NOT EXISTS
+        FOR (n:Entity) ON (n.name_embedding)
+        OPTIONS {indexConfig: {
+         `vector.dimensions`: 1024,
+         `vector.similarity_function`: 'cosine'
+        }}
+        """,
+    ]
+    index_queries: list[LiteralString] = range_indices + fulltext_indices + vector_indices
+
+    await asyncio.gather(*[driver.execute_query(query) for query in index_queries])
+
+
+async def clear_data(driver: AsyncDriver):
+    async with driver.session() as session:
+
+        async def delete_all(tx):
+            await tx.run('MATCH (n) DETACH DELETE n')
+
+        await session.execute_write(delete_all)
+
+
+async def retrieve_episodes(
+    driver: AsyncDriver,
+    reference_time: datetime,
+    last_n: int = EPISODE_WINDOW_LEN,
+) -> list[EpisodicNode]:
+    """
+    Retrieve the last n episodic nodes from the graph.
+
+    Args:
+        driver (AsyncDriver): The Neo4j driver instance.
+        reference_time (datetime): The reference time to filter episodes. Only episodes with a valid_at timestamp
+                                   less than or equal to this reference_time will be retrieved. This allows for
+                                   querying the graph's state at a specific point in time.
+        last_n (int, optional): The number of most recent episodes to retrieve, relative to the reference_time.
+
+    Returns:
+        list[EpisodicNode]: A list of EpisodicNode objects representing the retrieved episodes.
+    """
+    result = await driver.execute_query(
+        """
+        MATCH (e:Episodic) WHERE e.valid_at <= $reference_time
+        RETURN e.content as content,
+            e.created_at as created_at,
+            e.valid_at as valid_at,
+            e.uuid as uuid,
+            e.name as name,
+            e.source_description as source_description,
+            e.source as source
+        ORDER BY e.created_at DESC
+        LIMIT $num_episodes
+        """,
+        reference_time=reference_time,
+        num_episodes=last_n,
+    )
+    episodes = [
+        EpisodicNode(
+            content=record['content'],
+            created_at=datetime.fromtimestamp(
+                record['created_at'].to_native().timestamp(), timezone.utc
+            ),
+            valid_at=(record['valid_at'].to_native()),
+            uuid=record['uuid'],
+            source=EpisodeType.from_str(record['source']),
+            name=record['name'],
+            source_description=record['source_description'],
+        )
+        for record in result.records
+    ]
+    return list(reversed(episodes))  # Return in chronological order

graphiti_core/utils/maintenance/node_operations.py

@@ -0,0 +1,199 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+from datetime import datetime
+from time import time
+from typing import Any
+
+from graphiti_core.llm_client import LLMClient
+from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
+from graphiti_core.prompts import prompt_library
+
+logger = logging.getLogger(__name__)
+
+
+async def extract_message_nodes(
+    llm_client: LLMClient, episode: EpisodicNode, previous_episodes: list[EpisodicNode]
+) -> list[dict[str, Any]]:
+    # Prepare context for LLM
+    context = {
+        'episode_content': episode.content,
+        'episode_timestamp': episode.valid_at.isoformat(),
+        'previous_episodes': [
+            {
+                'content': ep.content,
+                'timestamp': ep.valid_at.isoformat(),
+            }
+            for ep in previous_episodes
+        ],
+    }
+
+    llm_response = await llm_client.generate_response(prompt_library.extract_nodes.v2(context))
+    extracted_node_data = llm_response.get('extracted_nodes', [])
+    return extracted_node_data
+
+
+async def extract_json_nodes(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+) -> list[dict[str, Any]]:
+    # Prepare context for LLM
+    context = {
+        'episode_content': episode.content,
+        'episode_timestamp': episode.valid_at.isoformat(),
+        'source_description': episode.source_description,
+    }
+
+    llm_response = await llm_client.generate_response(
+        prompt_library.extract_nodes.extract_json(context)
+    )
+    extracted_node_data = llm_response.get('extracted_nodes', [])
+    return extracted_node_data
+
+
+async def extract_nodes(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+    previous_episodes: list[EpisodicNode],
+) -> list[EntityNode]:
+    start = time()
+    extracted_node_data: list[dict[str, Any]] = []
+    if episode.source in [EpisodeType.message, EpisodeType.text]:
+        extracted_node_data = await extract_message_nodes(llm_client, episode, previous_episodes)
+    elif episode.source == EpisodeType.json:
+        extracted_node_data = await extract_json_nodes(llm_client, episode)
+
+    end = time()
+    logger.info(f'Extracted new nodes: {extracted_node_data} in {(end - start) * 1000} ms')
+    # Convert the extracted data into EntityNode objects
+    new_nodes = []
+    for node_data in extracted_node_data:
+        new_node = EntityNode(
+            name=node_data['name'],
+            labels=node_data['labels'],
+            summary=node_data['summary'],
+            created_at=datetime.now(),
+        )
+        new_nodes.append(new_node)
+        logger.info(f'Created new node: {new_node.name} (UUID: {new_node.uuid})')
+
+    return new_nodes
+
+
+async def dedupe_extracted_nodes(
+    llm_client: LLMClient,
+    extracted_nodes: list[EntityNode],
+    existing_nodes: list[EntityNode],
+) -> tuple[list[EntityNode], dict[str, str], list[EntityNode]]:
+    start = time()
+
+    # build existing node map
+    node_map: dict[str, EntityNode] = {}
+    for node in existing_nodes:
+        node_map[node.name] = node
+
+    # Temp hack
+    new_nodes_map: dict[str, EntityNode] = {}
+    for node in extracted_nodes:
+        new_nodes_map[node.name] = node
+
+    # Prepare context for LLM
+    existing_nodes_context = [
+        {'name': node.name, 'summary': node.summary} for node in existing_nodes
+    ]
+
+    extracted_nodes_context = [
+        {'name': node.name, 'summary': node.summary} for node in extracted_nodes
+    ]
+
+    context = {
+        'existing_nodes': existing_nodes_context,
+        'extracted_nodes': extracted_nodes_context,
+    }
+
+    llm_response = await llm_client.generate_response(prompt_library.dedupe_nodes.v2(context))
+
+    duplicate_data = llm_response.get('duplicates', [])
+
+    end = time()
+    logger.info(f'Deduplicated nodes: {duplicate_data} in {(end - start) * 1000} ms')
+
+    uuid_map: dict[str, str] = {}
+    for duplicate in duplicate_data:
+        uuid = new_nodes_map[duplicate['name']].uuid
+        uuid_value = node_map[duplicate['duplicate_of']].uuid
+        uuid_map[uuid] = uuid_value
+
+    nodes: list[EntityNode] = []
+    brand_new_nodes: list[EntityNode] = []
+    for node in extracted_nodes:
+        if node.uuid in uuid_map:
+            existing_uuid = uuid_map[node.uuid]
+            # TODO(Preston): This is a bit of a hack I implemented because we were getting incorrect uuids for existing nodes,
+            # can you revisit the node dedup function and make it somewhat cleaner and add more comments/tests please?
+            # find an existing node by the uuid from the nodes_map (each key is name, so we need to iterate by uuid value)
+            existing_node = next((v for k, v in node_map.items() if v.uuid == existing_uuid), None)
+            if existing_node:
+                nodes.append(existing_node)
+
+            continue
+        brand_new_nodes.append(node)
+        nodes.append(node)
+
+    return nodes, uuid_map, brand_new_nodes
+
+
+async def dedupe_node_list(
+    llm_client: LLMClient,
+    nodes: list[EntityNode],
+) -> tuple[list[EntityNode], dict[str, str]]:
+    start = time()
+
+    # build node map
+    node_map = {}
+    for node in nodes:
+        node_map[node.name] = node
+
+    # Prepare context for LLM
+    nodes_context = [{'name': node.name, 'summary': node.summary} for node in nodes]
+
+    context = {
+        'nodes': nodes_context,
+    }
+
+    llm_response = await llm_client.generate_response(
+        prompt_library.dedupe_nodes.node_list(context)
+    )
+
+    nodes_data = llm_response.get('nodes', [])
+
+    end = time()
+    logger.info(f'Deduplicated nodes: {nodes_data} in {(end - start) * 1000} ms')
+
+    # Get full node data
+    unique_nodes = []
+    uuid_map: dict[str, str] = {}
+    for node_data in nodes_data:
+        node = node_map[node_data['names'][0]]
+        unique_nodes.append(node)
+
+        for name in node_data['names'][1:]:
+            uuid = node_map[name].uuid
+            uuid_value = node_map[node_data['names'][0]].uuid
+            uuid_map[uuid] = uuid_value
+
+    return unique_nodes, uuid_map

graphiti_core/utils/maintenance/temporal_operations.py

@@ -0,0 +1,184 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+from datetime import datetime
+from typing import List
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.llm_client import LLMClient
+from graphiti_core.nodes import EntityNode, EpisodicNode
+from graphiti_core.prompts import prompt_library
+
+logger = logging.getLogger(__name__)
+
+NodeEdgeNodeTriplet = tuple[EntityNode, EntityEdge, EntityNode]
+
+
+def extract_node_and_edge_triplets(
+    edges: list[EntityEdge], nodes: list[EntityNode]
+) -> list[NodeEdgeNodeTriplet]:
+    return [extract_node_edge_node_triplet(edge, nodes) for edge in edges]
+
+
+def extract_node_edge_node_triplet(
+    edge: EntityEdge, nodes: list[EntityNode]
+) -> NodeEdgeNodeTriplet:
+    source_node = next((node for node in nodes if node.uuid == edge.source_node_uuid), None)
+    target_node = next((node for node in nodes if node.uuid == edge.target_node_uuid), None)
+    if not source_node or not target_node:
+        raise ValueError(f'Source or target node not found for edge {edge.uuid}')
+    return (source_node, edge, target_node)
+
+
+def prepare_edges_for_invalidation(
+    existing_edges: list[EntityEdge],
+    new_edges: list[EntityEdge],
+    nodes: list[EntityNode],
+) -> tuple[list[NodeEdgeNodeTriplet], list[NodeEdgeNodeTriplet]]:
+    existing_edges_pending_invalidation: list[NodeEdgeNodeTriplet] = []
+    new_edges_with_nodes: list[NodeEdgeNodeTriplet] = []
+
+    for edge_list, result_list in [
+        (existing_edges, existing_edges_pending_invalidation),
+        (new_edges, new_edges_with_nodes),
+    ]:
+        for edge in edge_list:
+            source_node = next((node for node in nodes if node.uuid == edge.source_node_uuid), None)
+            target_node = next((node for node in nodes if node.uuid == edge.target_node_uuid), None)
+
+            if source_node and target_node:
+                result_list.append((source_node, edge, target_node))
+
+    return existing_edges_pending_invalidation, new_edges_with_nodes
+
+
+async def invalidate_edges(
+    llm_client: LLMClient,
+    existing_edges_pending_invalidation: list[NodeEdgeNodeTriplet],
+    new_edges: list[NodeEdgeNodeTriplet],
+    current_episode: EpisodicNode,
+    previous_episodes: list[EpisodicNode],
+) -> list[EntityEdge]:
+    invalidated_edges = []  # TODO: this is not yet used?
+
+    context = prepare_invalidation_context(
+        existing_edges_pending_invalidation,
+        new_edges,
+        current_episode,
+        previous_episodes,
+    )
+    llm_response = await llm_client.generate_response(prompt_library.invalidate_edges.v1(context))
+
+    edges_to_invalidate = llm_response.get('invalidated_edges', [])
+    invalidated_edges = process_edge_invalidation_llm_response(
+        edges_to_invalidate, existing_edges_pending_invalidation
+    )
+
+    return invalidated_edges
+
+
+def extract_date_strings_from_edge(edge: EntityEdge) -> str:
+    start = edge.valid_at
+    end = edge.invalid_at
+    date_string = f'Start Date: {start.isoformat()}' if start else ''
+    if end:
+        date_string += f' (End Date: {end.isoformat()})'
+    return date_string
+
+
+def prepare_invalidation_context(
+    existing_edges: list[NodeEdgeNodeTriplet],
+    new_edges: list[NodeEdgeNodeTriplet],
+    current_episode: EpisodicNode,
+    previous_episodes: list[EpisodicNode],
+) -> dict:
+    return {
+        'existing_edges': [
+            f'{edge.uuid} | {source_node.name} - {edge.name} - {target_node.name} (Fact: {edge.fact}) {extract_date_strings_from_edge(edge)}'
+            for source_node, edge, target_node in sorted(
+                existing_edges, key=lambda x: (x[1].created_at), reverse=True
+            )
+        ],
+        'new_edges': [
+            f'{edge.uuid} | {source_node.name} - {edge.name} - {target_node.name} (Fact: {edge.fact}) {extract_date_strings_from_edge(edge)}'
+            for source_node, edge, target_node in sorted(
+                new_edges, key=lambda x: (x[1].created_at), reverse=True
+            )
+        ],
+        'current_episode': current_episode.content,
+        'previous_episodes': [episode.content for episode in previous_episodes],
+    }
+
+
+def process_edge_invalidation_llm_response(
+    edges_to_invalidate: List[dict], existing_edges: List[NodeEdgeNodeTriplet]
+) -> List[EntityEdge]:
+    invalidated_edges = []
+    for edge_to_invalidate in edges_to_invalidate:
+        edge_uuid = edge_to_invalidate['edge_uuid']
+        edge_to_update = next(
+            (edge for _, edge, _ in existing_edges if edge.uuid == edge_uuid),
+            None,
+        )
+        if edge_to_update:
+            edge_to_update.expired_at = datetime.now()
+            edge_to_update.fact = edge_to_invalidate['fact']
+            invalidated_edges.append(edge_to_update)
+            logger.info(
+                f"Invalidated edge: {edge_to_update.name} (UUID: {edge_to_update.uuid}). Updated Fact: {edge_to_invalidate['fact']}"
+            )
+    return invalidated_edges
+
+
+async def extract_edge_dates(
+    llm_client: LLMClient,
+    edge: EntityEdge,
+    reference_time: datetime,
+    current_episode: EpisodicNode,
+    previous_episodes: List[EpisodicNode],
+) -> tuple[datetime | None, datetime | None, str]:
+    context = {
+        'edge_name': edge.name,
+        'edge_fact': edge.fact,
+        'current_episode': current_episode.content,
+        'previous_episodes': [ep.content for ep in previous_episodes],
+        'reference_timestamp': reference_time.isoformat(),
+    }
+    llm_response = await llm_client.generate_response(prompt_library.extract_edge_dates.v1(context))
+
+    valid_at = llm_response.get('valid_at')
+    invalid_at = llm_response.get('invalid_at')
+    explanation = llm_response.get('explanation', '')
+
+    valid_at_datetime = None
+    invalid_at_datetime = None
+
+    if valid_at and valid_at != '':
+        try:
+            valid_at_datetime = datetime.fromisoformat(valid_at.replace('Z', '+00:00'))
+        except ValueError as e:
+            logger.error(f'Error parsing valid_at date: {e}. Input: {valid_at}')
+
+    if invalid_at and invalid_at != '':
+        try:
+            invalid_at_datetime = datetime.fromisoformat(invalid_at.replace('Z', '+00:00'))
+        except ValueError as e:
+            logger.error(f'Error parsing invalid_at date: {e}. Input: {invalid_at}')
+
+    logger.info(f'Edge date extraction explanation: {explanation}')
+
+    return valid_at_datetime, invalid_at_datetime, explanation