Skip to content

Commit 82c339a

Browse files
natoverseCopilot
andauthored
Sync docs for schema and metadata fields (#2435)
* Sync docs for schema and metadata fields * Clarify prepend_metadata raw_data note * Align docs notebooks with v3 model APIs * Clean notebook formatting in v3 model cells * Use clean global_search params without pop workaround * Fix dynamic global search notebook JSON mode param conflict * Remove unused imports in docs drift notebook * Restore drift notebook data setup variables * Apply ruff format to drift notebook * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * remove outputs --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent f5af4d2 commit 82c339a

9 files changed

Lines changed: 515 additions & 170 deletions

File tree

docs/config/yaml.md

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ embedding_models:
7474

7575
### input
7676

77-
Our pipeline can ingest .csv, .txt, or .json data from an input location. See the [inputs page](../index/inputs.md) for more details and examples.
77+
Our pipeline can ingest `.csv`, `.txt`, `.json`, `.jsonl`, `.parquet`, or MarkItDown-supported files from an input location. See the [inputs page](../index/inputs.md) for more details and examples.
7878

7979
#### Fields
8080

@@ -86,24 +86,24 @@ Our pipeline can ingest .csv, .txt, or .json data from an input location. See th
8686
- `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
8787
- `account_url` **str** - (blob only) The storage account blob URL to use.
8888
- `database_name` **str** - (cosmosdb only) The database name to use.
89-
- `type` **text|csv|json** - The type of input data to load. Default is `text`
89+
- `type` **text|csv|json|jsonl|parquet|markitdown** - The type of input data to load. Default is `text`
9090
- `encoding` **str** - The encoding of the input file. Default is `utf-8`
91-
- `file_pattern` **str** - A regex to match input files. Default is `.*\.csv$`, `.*\.txt$`, or `.*\.json$` depending on the specified `type`, but you can customize it if needed.
91+
- `file_pattern` **str** - A regex to match input files. Defaults are derived from `type` (for example `.*\.csv$`, `.*\.txt$`, `.*\.json$`, `.*\.jsonl$`, `.*\.parquet$`), but you can customize it if needed.
9292
- `id_column` **str** - The input ID column to use.
9393
- `title_column` **str** - The input title column to use.
9494
- `text_column` **str** - The input text column to use.
9595

9696
### chunking
9797

98-
These settings configure how we parse documents into text chunks. This is necessary because very large documents may not fit into a single context window, and graph extraction accuracy can be modulated. Also note the `metadata` setting in the input document config, which will replicate document metadata into each chunk.
98+
These settings configure how we parse documents into text chunks. This is necessary because very large documents may not fit into a single context window, and graph extraction accuracy can be modulated. You can also prepend selected document fields (standard fields like `title` or values from `raw_data`) into each chunk.
9999

100100
#### Fields
101101

102102
- `type` **tokens|sentence** - The chunking type to use.
103103
- `encoding_model` **str** - The text encoding model to use for splitting on token boundaries.
104104
- `size` **int** - The max chunk size in tokens.
105105
- `overlap` **int** - The chunk overlap in tokens.
106-
- `prepend_metadata` **list[str]** - Metadata fields from the source document to prepend on each chunk.
106+
- `prepend_metadata` **list[str]** - Document fields to prepend on each chunk. These can be standard document fields (`id`, `title`, `text`, `creation_date`) or fields from the source row/object stored in `raw_data`. For structured inputs, the `raw_data` object contains any other fields present in the source file (for CSV, this is all other column data).
107107

108108
## Outputs and Storage
109109

@@ -120,8 +120,6 @@ This section controls the storage mechanism used by the pipeline used for export
120120
- `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
121121
- `account_url` **str** - (blob only) The storage account blob URL to use.
122122
- `database_name` **str** - (cosmosdb only) The database name to use.
123-
- `type` **text|csv|json** - The type of input data to load. Default is `text`
124-
- `encoding` **str** - The encoding of the input file. Default is `utf-8`
125123

126124
### update_output_storage
127125

@@ -136,8 +134,6 @@ The section defines a secondary storage location for running incremental indexin
136134
- `container_name` **str** - (blob/cosmosdb only) The Azure Storage container name.
137135
- `account_url` **str** - (blob only) The storage account blob URL to use.
138136
- `database_name` **str** - (cosmosdb only) The database name to use.
139-
- `type` **text|csv|json** - The type of input data to load. Default is `text`
140-
- `encoding` **str** - The encoding of the input file. Default is `utf-8`
141137

142138
### cache
143139

docs/examples_notebooks/drift_search.ipynb

Lines changed: 50 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,10 @@
1919
"import os\n",
2020
"\n",
2121
"import pandas as pd\n",
22-
"from graphrag.config.enums import ModelType\n",
2322
"from graphrag.config.models.drift_search_config import DRIFTSearchConfig\n",
24-
"from graphrag.config.models.language_model_config import LanguageModelConfig\n",
25-
"from graphrag.language_model.manager import ModelManager\n",
2623
"from graphrag.query.indexer_adapters import (\n",
2724
" read_indexer_entities,\n",
2825
" read_indexer_relationships,\n",
29-
" read_indexer_report_embeddings,\n",
3026
" read_indexer_reports,\n",
3127
" read_indexer_text_units,\n",
3228
")\n",
@@ -35,60 +31,10 @@
3531
")\n",
3632
"from graphrag.query.structured_search.drift_search.search import DRIFTSearch\n",
3733
"from graphrag.tokenizer.get_tokenizer import get_tokenizer\n",
38-
"from graphrag_vectors.lancedb import LanceDBVectorStore\n",
39-
"\n",
40-
"INPUT_DIR = \"./inputs/operation dulce\"\n",
41-
"LANCEDB_URI = f\"{INPUT_DIR}/lancedb\"\n",
42-
"\n",
43-
"COMMUNITY_REPORT_TABLE = \"community_reports\"\n",
44-
"COMMUNITY_TABLE = \"communities\"\n",
45-
"ENTITY_TABLE = \"entities\"\n",
46-
"RELATIONSHIP_TABLE = \"relationships\"\n",
47-
"COVARIATE_TABLE = \"covariates\"\n",
48-
"TEXT_UNIT_TABLE = \"text_units\"\n",
49-
"COMMUNITY_LEVEL = 2\n",
50-
"\n",
51-
"\n",
52-
"# read nodes table to get community and degree data\n",
53-
"entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n",
54-
"community_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet\")\n",
55-
"\n",
56-
"print(f\"Entity df columns: {entity_df.columns}\")\n",
57-
"\n",
58-
"entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)\n",
59-
"\n",
60-
"# load description embeddings to an in-memory lancedb vectorstore\n",
61-
"# to connect to a remote db, specify url and port values.\n",
62-
"description_embedding_store = LanceDBVectorStore(\n",
63-
" db_uri=LANCEDB_URI,\n",
64-
" index_name=\"entity_description\",\n",
65-
")\n",
66-
"description_embedding_store.connect()\n",
67-
"\n",
68-
"full_content_embedding_store = LanceDBVectorStore(\n",
69-
" db_uri=LANCEDB_URI,\n",
70-
" index_name=\"community_full_content\",\n",
71-
")\n",
72-
"full_content_embedding_store.connect()\n",
73-
"\n",
74-
"print(f\"Entity count: {len(entity_df)}\")\n",
75-
"entity_df.head()\n",
76-
"\n",
77-
"relationship_df = pd.read_parquet(f\"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet\")\n",
78-
"relationships = read_indexer_relationships(relationship_df)\n",
79-
"\n",
80-
"print(f\"Relationship count: {len(relationship_df)}\")\n",
81-
"relationship_df.head()\n",
82-
"\n",
83-
"text_unit_df = pd.read_parquet(f\"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet\")\n",
84-
"text_units = read_indexer_text_units(text_unit_df)\n",
85-
"\n",
86-
"print(f\"Text unit records: {len(text_unit_df)}\")\n",
87-
"text_unit_df.head()\n",
88-
"\n",
89-
"report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n",
90-
"reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)\n",
91-
"read_indexer_report_embeddings(reports, full_content_embedding_store)"
34+
"from graphrag_llm.completion import create_completion\n",
35+
"from graphrag_llm.config import ModelConfig\n",
36+
"from graphrag_llm.embedding import create_embedding\n",
37+
"from graphrag_vectors import IndexSchema, LanceDBVectorStore"
9238
]
9339
},
9440
{
@@ -99,33 +45,61 @@
9945
"source": [
10046
"api_key = os.environ[\"GRAPHRAG_API_KEY\"]\n",
10147
"\n",
102-
"chat_config = LanguageModelConfig(\n",
103-
" api_key=api_key,\n",
104-
" type=ModelType.Chat,\n",
48+
"chat_config = ModelConfig(\n",
49+
" type=\"litellm\",\n",
10550
" model_provider=\"openai\",\n",
10651
" model=\"gpt-4.1\",\n",
107-
" max_retries=20,\n",
108-
")\n",
109-
"chat_model = ModelManager().get_or_create_chat_model(\n",
110-
" name=\"local_search\",\n",
111-
" model_type=ModelType.Chat,\n",
112-
" config=chat_config,\n",
52+
" api_key=api_key,\n",
11353
")\n",
54+
"chat_model = create_completion(chat_config)\n",
11455
"\n",
11556
"tokenizer = get_tokenizer(chat_config)\n",
11657
"\n",
117-
"embedding_config = LanguageModelConfig(\n",
118-
" api_key=api_key,\n",
119-
" type=ModelType.Embedding,\n",
58+
"embedding_config = ModelConfig(\n",
59+
" type=\"litellm\",\n",
12060
" model_provider=\"openai\",\n",
12161
" model=\"text-embedding-3-large\",\n",
122-
" max_retries=20,\n",
62+
" api_key=api_key,\n",
63+
")\n",
64+
"\n",
65+
"text_embedder = create_embedding(embedding_config)"
66+
]
67+
},
68+
{
69+
"cell_type": "code",
70+
"execution_count": null,
71+
"metadata": {},
72+
"outputs": [],
73+
"source": [
74+
"# parquet files generated from indexing pipeline\n",
75+
"INPUT_DIR = \"./inputs/operation dulce\"\n",
76+
"LANCEDB_URI = \"./lancedb\"\n",
77+
"COMMUNITY_TABLE = \"communities\"\n",
78+
"COMMUNITY_REPORT_TABLE = \"community_reports\"\n",
79+
"ENTITY_TABLE = \"entities\"\n",
80+
"RELATIONSHIP_TABLE = \"relationships\"\n",
81+
"TEXT_UNIT_TABLE = \"text_units\"\n",
82+
"COMMUNITY_LEVEL = 2\n",
83+
"\n",
84+
"community_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet\")\n",
85+
"report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n",
86+
"entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n",
87+
"relationship_df = pd.read_parquet(f\"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet\")\n",
88+
"text_unit_df = pd.read_parquet(f\"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet\")\n",
89+
"\n",
90+
"reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)\n",
91+
"entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)\n",
92+
"relationships = read_indexer_relationships(relationship_df)\n",
93+
"text_units = read_indexer_text_units(text_unit_df)\n",
94+
"\n",
95+
"# Connect to the existing GraphRAG vector index for entity-description embeddings.\n",
96+
"description_embedding_store = LanceDBVectorStore(\n",
97+
" index_schema=IndexSchema(index_name=\"default-entity-description\")\n",
12398
")\n",
99+
"description_embedding_store.connect(db_uri=LANCEDB_URI)\n",
124100
"\n",
125-
"text_embedder = ModelManager().get_or_create_embedding_model(\n",
126-
" name=\"local_search_embedding\",\n",
127-
" model_type=ModelType.Embedding,\n",
128-
" config=embedding_config,\n",
101+
"print(\n",
102+
" f\"Loaded reports={len(reports)}, entities={len(entities)}, relationships={len(relationships)}, text_units={len(text_units)}\"\n",
129103
")"
130104
]
131105
},
@@ -188,7 +162,7 @@
188162
],
189163
"metadata": {
190164
"kernelspec": {
191-
"display_name": "Python 3",
165+
"display_name": "graphrag-monorepo (3.12.10)",
192166
"language": "python",
193167
"name": "python3"
194168
},

0 commit comments

Comments
 (0)