Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions examples/docs_to_knowledge_graph/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,9 +82,15 @@ def docs_to_kg_flow(
cocoindex.functions.ExtractByLlm(
llm_spec=cocoindex.LlmSpec(
# Supported LLM: https://cocoindex.io/docs/ai/llm
api_type=cocoindex.LlmApiType.OPENAI,
model="gpt-4o",
api_type=cocoindex.LlmApiType.OLLAMA,
model="llama3.2",
address="http://localhost:11434",
),
# Alternative: Use OpenAI API model instead of Ollama
# llm_spec=cocoindex.LlmSpec(
# api_type=cocoindex.LlmApiType.OPENAI,
# model="gpt-4o",
# ),
output_type=DocumentSummary,
instruction="Please summarize the content of the document.",
)
Expand All @@ -100,9 +106,15 @@ def docs_to_kg_flow(
cocoindex.functions.ExtractByLlm(
llm_spec=cocoindex.LlmSpec(
# Supported LLM: https://cocoindex.io/docs/ai/llm
api_type=cocoindex.LlmApiType.OPENAI,
model="gpt-4o",
api_type=cocoindex.LlmApiType.OLLAMA,
model="llama3.2",
address="http://localhost:11434",
),
# Alternative: Use OpenAI API model instead of Ollama
# llm_spec=cocoindex.LlmSpec(
# api_type=cocoindex.LlmApiType.OPENAI,
# model="gpt-4o",
# ),
output_type=list[Relationship],
instruction=(
"Please extract relationships from CocoIndex documents. "
Expand Down
2 changes: 1 addition & 1 deletion src/base/spec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -384,7 +384,7 @@ impl fmt::Display for VectorSimilarityMetric {
}
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind")]
pub enum VectorIndexMethod {
Hnsw {
Expand Down
37 changes: 31 additions & 6 deletions src/ops/targets/neo4j.rs
Original file line number Diff line number Diff line change
Expand Up @@ -568,9 +568,6 @@ impl SetupState {
.map(|f| (f.name.as_str(), &f.value_type.typ))
.collect::<HashMap<_, _>>();
for index_def in index_options.vector_indexes.iter() {
if index_def.method.is_some() {
api_bail!("Vector index method is not configurable for Neo4j yet");
}
sub_components.push(ComponentState {
object_label: schema.elem_type.clone(),
index_def: IndexDef::from_vector_index_def(
Expand All @@ -583,6 +580,7 @@ impl SetupState {
index_def.field_name
)
})?,
index_def.method.clone(),
)?,
});
}
Expand Down Expand Up @@ -644,14 +642,20 @@ enum IndexDef {
field_name: String,
metric: spec::VectorSimilarityMetric,
vector_size: usize,
method: Option<spec::VectorIndexMethod>,
},
}

impl IndexDef {
fn from_vector_index_def(
index_def: &spec::VectorIndexDef,
field_typ: &schema::ValueType,
_method: Option<spec::VectorIndexMethod>,
) -> Result<Self> {
let method = index_def.method.clone();
if let Some(spec::VectorIndexMethod::IvfFlat { .. }) = method {
api_bail!("IVFFlat vector index method is not supported for Neo4j");
}
Ok(Self::VectorIndex {
field_name: index_def.field_name.clone(),
vector_size: (match field_typ {
Expand All @@ -664,6 +668,7 @@ impl IndexDef {
api_error!("Vector index field must be a vector with fixed dimension")
})?,
metric: index_def.metric,
method,
})
}
}
Expand Down Expand Up @@ -723,9 +728,14 @@ impl components::SetupOperator for SetupComponentOperator {
field_name,
metric,
vector_size,
method,
} => {
let method_str = method
.as_ref()
.map(|m| format!(", method: {}", m))
.unwrap_or_default();
format!(
"{key_desc} ON {label} (field_name: {field_name}, vector_size: {vector_size}, metric: {metric})",
"{key_desc} ON {label} (field_name: {field_name}, vector_size: {vector_size}, metric: {metric}{method_str})",
)
}
}
Expand All @@ -752,17 +762,32 @@ impl components::SetupOperator for SetupComponentOperator {
field_name,
metric,
vector_size,
method,
} => {
let mut parts = vec![];

parts.push(format!("`vector.dimensions`: {}", vector_size));
parts.push(format!("`vector.similarity_function`: '{}'", metric));

if let Some(spec::VectorIndexMethod::Hnsw { m, ef_construction }) = method {
if let Some(m_val) = m {
parts.push(format!("`vector.hnsw.m`: {}", m_val));
}
if let Some(ef_val) = ef_construction {
parts.push(format!("`vector.hnsw.ef_construction`: {}", ef_val));
}
}

formatdoc! {"
CREATE VECTOR INDEX {name} IF NOT EXISTS
FOR {matcher} ON {qualifier}.{field_name}
OPTIONS {{
indexConfig: {{
`vector.dimensions`: {vector_size},
`vector.similarity_function`: '{metric}'
{config}
}}
}}",
name = key.name,
config = parts.join(", ")
}
}
});
Expand Down