|
| 1 | +""" |
| 2 | +This example shows how to extract relationships from Markdown documents and build a knowledge graph. |
| 3 | +""" |
| 4 | +import dataclasses |
| 5 | +from dotenv import load_dotenv |
| 6 | +import cocoindex |
| 7 | + |
| 8 | + |
| 9 | +@dataclasses.dataclass |
| 10 | +class Relationship: |
| 11 | + """Describe a relationship between two nodes.""" |
| 12 | + subject: str |
| 13 | + predicate: str |
| 14 | + object: str |
| 15 | + |
| 16 | +@dataclasses.dataclass |
| 17 | +class Relationships: |
| 18 | + """Describe a relationship between two nodes.""" |
| 19 | + relationships: list[Relationship] |
| 20 | + |
| 21 | +@cocoindex.flow_def(name="DocsToKG") |
| 22 | +def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope): |
| 23 | + """ |
| 24 | + Define an example flow that extracts triples from files and build knowledge graph. |
| 25 | + """ |
| 26 | + |
| 27 | + conn_spec = cocoindex.add_auth_entry( |
| 28 | + "Neo4jConnection", |
| 29 | + cocoindex.storages.Neo4jConnectionSpec( |
| 30 | + uri="bolt://localhost:7687", |
| 31 | + user="neo4j", |
| 32 | + password="cocoindex", |
| 33 | + )) |
| 34 | + |
| 35 | + data_scope["documents"] = flow_builder.add_source( |
| 36 | + cocoindex.sources.LocalFile(path="../../docs/docs/core", |
| 37 | + included_patterns=["*.md", "*.mdx"])) |
| 38 | + |
| 39 | + relationships = data_scope.add_collector() |
| 40 | + |
| 41 | + with data_scope["documents"].row() as doc: |
| 42 | + doc["chunks"] = doc["content"].transform( |
| 43 | + cocoindex.functions.SplitRecursively(), |
| 44 | + language="markdown", chunk_size=10000) |
| 45 | + |
| 46 | + with doc["chunks"].row() as chunk: |
| 47 | + chunk["relationships"] = chunk["text"].transform( |
| 48 | + cocoindex.functions.ExtractByLlm( |
| 49 | + llm_spec=cocoindex.LlmSpec( |
| 50 | + api_type=cocoindex.LlmApiType.OPENAI, model="gpt-4o"), |
| 51 | + output_type=Relationships, |
| 52 | + instruction=( |
| 53 | + "Please extract relationships from CocoIndex documents. " |
| 54 | + "Focus on concepts and ingnore specific examples. " |
| 55 | + "Each relationship should be a tuple of (subject, predicate, object)."))) |
| 56 | + |
| 57 | + with chunk["relationships"]["relationships"].row() as relationship: |
| 58 | + relationships.collect( |
| 59 | + id=cocoindex.GeneratedField.UUID, |
| 60 | + subject=relationship["subject"], |
| 61 | + predicate=relationship["predicate"], |
| 62 | + object=relationship["object"], |
| 63 | + ) |
| 64 | + |
| 65 | + relationships.export( |
| 66 | + "relationships", |
| 67 | + cocoindex.storages.Neo4jRelationship( |
| 68 | + connection=conn_spec, |
| 69 | + rel_type="RELATIONSHIP", |
| 70 | + source=cocoindex.storages.Neo4jRelationshipEndSpec( |
| 71 | + label="Entity", |
| 72 | + fields=[cocoindex.storages.Neo4jFieldMapping(field_name="subject", node_field_name="value")] |
| 73 | + ), |
| 74 | + target=cocoindex.storages.Neo4jRelationshipEndSpec( |
| 75 | + label="Entity", |
| 76 | + fields=[cocoindex.storages.Neo4jFieldMapping(field_name="object", node_field_name="value")] |
| 77 | + ), |
| 78 | + nodes={ |
| 79 | + "Entity": cocoindex.storages.Neo4jRelationshipNodeSpec(key_field_name="value"), |
| 80 | + }, |
| 81 | + ), |
| 82 | + primary_key_fields=["id"], |
| 83 | + ) |
| 84 | + |
| 85 | +@cocoindex.main_fn() |
| 86 | +def _run(): |
| 87 | + pass |
| 88 | + |
| 89 | +if __name__ == "__main__": |
| 90 | + load_dotenv(override=True) |
| 91 | + _run() |
0 commit comments