diff --git a/python/pyproject.toml b/python/pyproject.toml index ee735bdb1..85a4ef50d 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -7,6 +7,8 @@ dependencies = [ "fastapi>=0.104.0", "uvicorn>=0.24.0", "pydantic>=2.0.0", + "openai>=1.52.0", + "lance>=0.17.0", ] description = "Python bindings for the lance-graph Cypher engine" authors = [{ name = "Lance Devs", email = "dev@lancedb.com" }] @@ -43,8 +45,6 @@ build-backend = "maturin" [project.optional-dependencies] tests = ["pytest", "pyarrow>=14", "pandas", "ruff"] dev = ["ruff", "pyright"] -llm = ["openai>=1.52.0"] -lance-storage = ["lance>=0.17.0"] [project.scripts] knowledge_graph = "knowledge_graph.main:main" diff --git a/python/python/knowledge_graph/cli/__init__.py b/python/python/knowledge_graph/cli/__init__.py new file mode 100644 index 000000000..d65c91125 --- /dev/null +++ b/python/python/knowledge_graph/cli/__init__.py @@ -0,0 +1,9 @@ +from .ingest import extract_and_add, preview_extraction +from .interactive import list_datasets, run_interactive + +__all__ = [ + "run_interactive", + "list_datasets", + "preview_extraction", + "extract_and_add", +] diff --git a/python/python/knowledge_graph/cli/embedding_support.py b/python/python/knowledge_graph/cli/embedding_support.py new file mode 100644 index 000000000..77db3767a --- /dev/null +++ b/python/python/knowledge_graph/cli/embedding_support.py @@ -0,0 +1,19 @@ +"""Shared helpers for preparing rows and embeddings (internal).""" + +from __future__ import annotations + +from .ingest import ( + _assign_embeddings, + _format_entity_embedding_input, + _format_relationship_embedding_input, + _prepare_entity_rows, + _prepare_relationship_rows, +) + +__all__ = [ + "_assign_embeddings", + "_format_entity_embedding_input", + "_format_relationship_embedding_input", + "_prepare_entity_rows", + "_prepare_relationship_rows", +] diff --git a/python/python/knowledge_graph/cli/ingest.py b/python/python/knowledge_graph/cli/ingest.py new file mode 100644 index 000000000..e41ec73b2 --- /dev/null +++ b/python/python/knowledge_graph/cli/ingest.py @@ -0,0 +1,220 @@ +"""Extraction preview and ingest helpers for the knowledge graph CLI.""" + +from __future__ import annotations + +import hashlib +import json +import logging +from dataclasses import asdict, is_dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Mapping + +from .. import extraction as kg_extraction + +if TYPE_CHECKING: + from ..embeddings import EmbeddingGenerator + from ..service import LanceKnowledgeGraph + +LOGGER = logging.getLogger(__name__) + + +def preview_extraction(source: str, extractor: kg_extraction.BaseExtractor) -> None: + """Preview extracted knowledge from a text source or inline text.""" + text = _resolve_text_input(source) + result = kg_extraction.preview_extraction(text, extractor=extractor) + print(json.dumps(_result_to_dict(result), indent=2)) + + +def extract_and_add( + source: str, + service: LanceKnowledgeGraph, + extractor: kg_extraction.BaseExtractor, + *, + embedding_generator: EmbeddingGenerator | None = None, +) -> None: + """Extract knowledge and append it to the backing graph.""" + import pyarrow as pa + + text = _resolve_text_input(source) + result = kg_extraction.preview_extraction(text, extractor=extractor) + entity_rows, name_to_id = _prepare_entity_rows( + result.entities, embedding_generator=embedding_generator + ) + relationships = result.relationships + + if not entity_rows and not relationships: + print("No candidate entities or relationships detected.") + return + + if entity_rows: + entity_table = pa.Table.from_pylist(entity_rows) + service.upsert_table("Entity", entity_table, merge=True) + message = f"Upserted {entity_table.num_rows} entity rows into dataset 'Entity'." + print(message) + + relationship_rows = _prepare_relationship_rows( + relationships, + name_to_id, + embedding_generator=embedding_generator, + ) + if relationship_rows: + rel_table = pa.Table.from_pylist(relationship_rows) + service.upsert_table("RELATIONSHIP", rel_table, merge=True) + message = ( + "Upserted " + f"{rel_table.num_rows} relationship rows into dataset " + "'RELATIONSHIP'." + ) + print(message) + + +def _resolve_text_input(raw: str) -> str: + """Load text from a file if it exists, otherwise treat the string as content.""" + candidate = Path(raw) + if candidate.exists(): + if candidate.is_dir(): + raise IsADirectoryError(f"Expected text file, got directory: {candidate}") + return candidate.read_text(encoding="utf-8") + return raw + + +def _ensure_dict(item: object) -> dict: + if is_dataclass(item): + return asdict(item) # type: ignore[arg-type] + if isinstance(item, dict): + return item + raise TypeError(f"Unsupported extraction item type: {type(item)!r}") + + +def _result_to_dict(result: "kg_extraction.ExtractionResult") -> dict[str, list[dict]]: + return { + "entities": [asdict(entity) for entity in result.entities], + "relationships": [asdict(rel) for rel in result.relationships], + } + + +def _prepare_entity_rows( + entities: list[Any], + *, + embedding_generator: EmbeddingGenerator | None = None, +) -> tuple[list[dict[str, Any]], dict[str, str]]: + rows: list[dict[str, Any]] = [] + name_to_id: dict[str, str] = {} + for entity in entities: + payload = _ensure_dict(entity) + name = str(payload.get("name", "")).strip() + entity_type = str( + payload.get("entity_type") or payload.get("type") or "" + ).strip() + if not name: + continue + base = f"{name}|{entity_type}".encode("utf-8") + entity_id = hashlib.md5(base).hexdigest() + payload["entity_id"] = entity_id + payload["entity_type"] = entity_type or "UNKNOWN" + payload["name_lower"] = name.lower() + rows.append(payload) + name_to_id.setdefault(name.lower(), entity_id) + if embedding_generator and rows: + _assign_embeddings( + rows, + embedding_generator, + _format_entity_embedding_input, + ) + return rows, name_to_id + + +def _prepare_relationship_rows( + relationships: list[Any], + name_to_id: dict[str, str], + *, + embedding_generator: EmbeddingGenerator | None = None, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for relation in relationships: + payload = _ensure_dict(relation) + source_name = str( + payload.get("source_entity_name") or payload.get("source") or "" + ).strip() + target_name = str( + payload.get("target_entity_name") or payload.get("target") or "" + ).strip() + source_id = name_to_id.get(source_name.lower()) + target_id = name_to_id.get(target_name.lower()) + if not (source_id and target_id): + continue + payload["source_entity_id"] = source_id + payload["target_entity_id"] = target_id + payload["relationship_type"] = ( + payload.get("relationship_type") or payload.get("type") or "RELATED_TO" + ) + payload.setdefault("source_entity_name", source_name) + payload.setdefault("target_entity_name", target_name) + rows.append(payload) + if embedding_generator and rows: + _assign_embeddings( + rows, + embedding_generator, + _format_relationship_embedding_input, + ) + return rows + + +def _assign_embeddings( + rows: list[dict[str, Any]], + embedding_generator: EmbeddingGenerator, + formatter: Callable[[Mapping[str, Any]], str], +) -> None: + texts: list[str] = [] + indices: list[int] = [] + for idx, row in enumerate(rows): + text = formatter(row) + if text: + texts.append(text) + indices.append(idx) + if not texts: + return + try: + vectors = embedding_generator.embed(texts) + except Exception as exc: # pragma: no cover - defensive logging path + LOGGER.warning("Failed to generate embeddings: %s", exc) + return + if len(vectors) != len(indices): + LOGGER.warning( + "Mismatch between embedding count and row count: expected %s, got %s", + len(indices), + len(vectors), + ) + return + for idx, vector in zip(indices, vectors): + rows[idx]["embedding"] = vector + + +def _format_entity_embedding_input(row: Mapping[str, Any]) -> str: + name = str(row.get("name", "")).strip() + entity_type = str(row.get("entity_type", "")).strip() + context = str(row.get("context", "")).strip() + pieces = [] + if name: + pieces.append(name) + if entity_type: + pieces.append(f"Type: {entity_type}") + if context: + pieces.append(f"Context: {context}") + return " | ".join(pieces) + + +def _format_relationship_embedding_input(row: Mapping[str, Any]) -> str: + source = str(row.get("source_entity_name") or row.get("source") or "").strip() + target = str(row.get("target_entity_name") or row.get("target") or "").strip() + relationship_type = str(row.get("relationship_type", "")).strip() + description = str(row.get("description", "")).strip() + core: list[str] = [] + if source or target: + if relationship_type: + core.append(f"{source} -[{relationship_type}]-> {target}".strip()) + else: + core.append(f"{source} -> {target}".strip()) + if description: + core.append(f"Description: {description}") + return " | ".join(part for part in core if part) diff --git a/python/python/knowledge_graph/cli/interactive.py b/python/python/knowledge_graph/cli/interactive.py new file mode 100644 index 000000000..b6c84cc5f --- /dev/null +++ b/python/python/knowledge_graph/cli/interactive.py @@ -0,0 +1,117 @@ +"""Interactive shell and CLI display helpers for the knowledge graph.""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +from ..store import LanceGraphStore + +if TYPE_CHECKING: + from ..config import KnowledgeGraphConfig + from ..service import LanceKnowledgeGraph + +if TYPE_CHECKING: + import pyarrow as pa + + +def list_datasets(config: "KnowledgeGraphConfig") -> None: + """List the Lance datasets available under the configured root.""" + store = LanceGraphStore(config) + store.ensure_layout() + datasets = store.list_datasets() + if not datasets: + print("No Lance datasets found. Load data or run extraction first.") + return + print("Available Lance datasets:") + for name, path in sorted(datasets.items()): + print(f" - {name}: {path}") + + +def run_interactive(service: "LanceKnowledgeGraph") -> None: + """Enter an interactive shell for issuing Cypher queries.""" + print("Lance Knowledge Graph interactive shell") + print("Type ':help' for commands, or 'quit' to exit.") + + while True: + try: + text = input("kg> ").strip() + except EOFError: + print() + break + + if not text: + continue + lowered = text.lower() + if lowered in {"quit", "exit", "q"}: + break + if text.startswith(":"): + _handle_command(text, service) + continue + + _execute_query(service, text) + + +def _handle_command(command: str, service: "LanceKnowledgeGraph") -> None: + """Handle meta-commands in the interactive shell.""" + cmd = command.strip() + if cmd in {":help", ":h"}: + print("Commands:") + print(" :help Show this message") + print(" :datasets List persisted Lance datasets") + print(" :config Show the configured node/relationship mappings") + print(" quit/exit/q Leave the shell") + return + if cmd in {":datasets", ":ls"}: + list_datasets(service.store.config) + return + if cmd in {":config", ":schema"}: + _print_config_summary(service) + return + print(f"Unknown command: {command}") + + +def _print_config_summary(service: "LanceKnowledgeGraph") -> None: + """Print a brief summary of the graph configuration.""" + config = service.config + # GraphConfig does not currently expose direct iterators; rely on repr. + print("Graph configuration:") + print(f" {config!r}") + + +def _execute_query(service: "LanceKnowledgeGraph", statement: str) -> None: + """Execute a single Cypher statement and print results.""" + try: + result = service.run(statement) + except Exception as exc: # pragma: no cover - CLI feedback path + print(f"Query failed: {exc}", file=sys.stderr) + return + + _print_table(result) + + +def _print_table(table: "pa.Table") -> None: + """Render a PyArrow table in a simple textual format.""" + if table.num_rows == 0: + print("(no rows)") + return + + column_names = table.column_names + columns = [table.column(i).to_pylist() for i in range(len(column_names))] + widths = [] + for name, values in zip(column_names, columns): + str_values = ["" if value is None else str(value) for value in values] + if str_values: + width = max(len(name), *(len(value) for value in str_values)) + else: + width = len(name) + widths.append(width) + + header = " | ".join(name.ljust(width) for name, width in zip(column_names, widths)) + separator = "-+-".join("-" * width for width in widths) + print(header) + print(separator) + for row_values in zip(*columns): + str_row = ["" if value is None else str(value) for value in row_values] + line = " | ".join(value.ljust(width) for value, width in zip(str_row, widths)) + print(line) diff --git a/python/python/knowledge_graph/cli/parser.py b/python/python/knowledge_graph/cli/parser.py new file mode 100644 index 000000000..9ce671c9f --- /dev/null +++ b/python/python/knowledge_graph/cli/parser.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +from .. import extraction as kg_extraction +from ..embeddings import DEFAULT_EMBEDDING_MODEL + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="knowledge_graph", + description="Operate the Lance-backed knowledge graph.", + ) + parser.add_argument( + "--root", + type=Path, + help="Root directory for Lance datasets (default: ./knowledge_graph_data).", + ) + parser.add_argument( + "--schema", + type=Path, + help="Path to a YAML file describing node and relationship mappings.", + ) + parser.add_argument( + "--list-datasets", + action="store_true", + help="List stored Lance datasets and exit.", + ) + parser.add_argument( + "--extractor", + choices=["heuristic", "llm"], + default=kg_extraction.DEFAULT_STRATEGY, + help="Extraction strategy to use (default: llm).", + ) + parser.add_argument( + "--llm-model", + default="gpt-4o-mini", + help="LLM model identifier when using --extractor llm.", + ) + parser.add_argument( + "--llm-temperature", + type=float, + default=0.2, + help="Sampling temperature for --extractor llm (default: 0.2).", + ) + parser.add_argument( + "--llm-config", + type=Path, + help=( + "Optional YAML file with OpenAI client options (api_key, base_url, " + "headers, etc)." + ), + ) + parser.add_argument( + "--embedding-model", + default=DEFAULT_EMBEDDING_MODEL, + help=("OpenAI embedding model for semantic search (set to 'none' to disable)."), + ) + parser.add_argument( + "--seed-count", + type=int, + default=5, + help=( + "Maximum number of seed entities to surface from similarity search " + "(default: 5)." + ), + ) + parser.add_argument( + "--log-level", + default="WARNING", + help="Logging level (DEBUG, INFO, WARNING, ERROR).", + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument( + "--init", + action="store_true", + help="Initialize the knowledge graph storage.", + ) + group.add_argument( + "--extract-preview", + metavar="INPUT", + help="Preview extracted knowledge from a file path or raw text.", + ) + group.add_argument( + "--extract-and-add", + metavar="INPUT", + help="Extract knowledge from a file path or raw text and insert it.", + ) + group.add_argument( + "--ask", + metavar="QUESTION", + help="Ask a natural-language question over the knowledge graph.", + ) + parser.add_argument( + "query", + nargs="?", + help="Execute a single Cypher query against the Lance datasets.", + ) + return parser diff --git a/python/python/knowledge_graph/cli/runtime.py b/python/python/knowledge_graph/cli/runtime.py new file mode 100644 index 000000000..eabd1cd76 --- /dev/null +++ b/python/python/knowledge_graph/cli/runtime.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Optional + +import yaml + +from ..config import KnowledgeGraphConfig +from ..service import LanceKnowledgeGraph +from ..store import LanceGraphStore + + +def load_config(args) -> KnowledgeGraphConfig: + if args.root: + config = KnowledgeGraphConfig.from_root(Path(args.root)) + else: + config = KnowledgeGraphConfig.default() + if args.schema: + config = config.with_schema(Path(args.schema)) + return config + + +def load_service(config: KnowledgeGraphConfig) -> LanceKnowledgeGraph: + graph_config = config.load_graph_config() + storage = LanceGraphStore(config) + service = LanceKnowledgeGraph(graph_config, storage=storage) + service.ensure_initialized() + return service + + +def configure_logging(level: str) -> None: + normalized = level.upper() + numeric = getattr(logging, normalized, None) + if not isinstance(numeric, int): + raise ValueError(f"Invalid log level: {level}") + logging.basicConfig(level=numeric) + + +def load_llm_options(path: Optional[Path]) -> dict: + if not path: + return {} + if not path.exists(): + raise FileNotFoundError(f"LLM config file not found: {path}") + with path.open("r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) or {} + if not isinstance(data, dict): + raise ValueError("LLM config must be a mapping") + if "default_headers" not in data and "http_headers" in data: + headers = data.pop("http_headers") + if isinstance(headers, dict): + data["default_headers"] = headers + return data diff --git a/python/python/knowledge_graph/llm/__init__.py b/python/python/knowledge_graph/llm/__init__.py new file mode 100644 index 000000000..14ec73fdb --- /dev/null +++ b/python/python/knowledge_graph/llm/__init__.py @@ -0,0 +1,28 @@ +from .llm_utils import ( + create_llm_client, + load_llm_options, + resolve_embedding_generator, +) +from .prompts import ( + build_answer_prompt, + build_query_prompt, + build_type_hint_lines, + summarize_schema, +) +from .qa import ask_question, execute_queries, extract_query_plan +from .semantic import collect_seed_neighbors, find_seed_entities + +__all__ = [ + "ask_question", + "extract_query_plan", + "execute_queries", + "build_answer_prompt", + "build_query_prompt", + "build_type_hint_lines", + "summarize_schema", + "collect_seed_neighbors", + "find_seed_entities", + "create_llm_client", + "load_llm_options", + "resolve_embedding_generator", +] diff --git a/python/python/knowledge_graph/llm/llm_utils.py b/python/python/knowledge_graph/llm/llm_utils.py new file mode 100644 index 000000000..ded160c57 --- /dev/null +++ b/python/python/knowledge_graph/llm/llm_utils.py @@ -0,0 +1,62 @@ +"""LLM and embedding helpers for the knowledge graph CLI and QA.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Mapping, Optional + +import yaml + +from .. import extraction as kg_extraction +from ..embeddings import EmbeddingGenerator + +if TYPE_CHECKING: + from pathlib import Path + +LOGGER = logging.getLogger(__name__) + + +def load_llm_options(path: Optional["Path"]) -> dict: + if not path: + return {} + if not path.exists(): + raise FileNotFoundError(f"LLM config file not found: {path}") + with path.open("r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) or {} + if not isinstance(data, dict): + raise ValueError("LLM config must be a mapping") + if "default_headers" not in data and "http_headers" in data: + headers = data.pop("http_headers") + if isinstance(headers, dict): + data["default_headers"] = headers + return data + + +def create_llm_client( + *, + llm_model: str, + llm_temperature: float, + llm_options: Optional[Mapping[str, Any]] = None, +) -> kg_extraction.LLMClient: + resolved_options = dict(llm_options or {}) + return kg_extraction.get_llm_client( + llm_model=llm_model, + llm_temperature=llm_temperature, + llm_options=resolved_options, + ) + + +def resolve_embedding_generator( + *, + model_name: str | None, + options: Optional[Mapping[str, Any]] = None, +) -> EmbeddingGenerator | None: + model = (model_name or "").strip() + if not model or model.lower() == "none": + return None + client_options = dict(options or {}) + try: + return EmbeddingGenerator(model=model, client_options=client_options) + except RuntimeError as exc: + LOGGER.warning("Embeddings disabled: %s", exc) + return None diff --git a/python/python/knowledge_graph/llm/prompts.py b/python/python/knowledge_graph/llm/prompts.py new file mode 100644 index 000000000..aa317727d --- /dev/null +++ b/python/python/knowledge_graph/llm/prompts.py @@ -0,0 +1,298 @@ +"""Prompt and summary builders for LLM-assisted QA over the graph.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Mapping, Sequence + +if TYPE_CHECKING: + from ..service import LanceKnowledgeGraph + + +def summarize_schema( + service: "LanceKnowledgeGraph", + max_columns: int = 20, + max_value_samples: int = 10, +) -> str: + type_hints = service.store.config.type_hints() + lines = [] + for name in service.dataset_names(): + try: + table = service.load_table(name) + except Exception: + continue + if hasattr(table, "schema"): + columns = list(getattr(table.schema, "names", [])) + else: + columns = [] + if len(columns) > max_columns: + columns = columns[:max_columns] + ["..."] + column_summary = f"- {name}: {', '.join(columns)}" + lines.append(column_summary) + + extras: list[str] = [] + upper_name = name.upper() + if upper_name == "ENTITY" and type_hints.get("entity_types"): + allowed = ", ".join(type_hints["entity_types"]) + extras.append(f"allowed entity_type values: {allowed}") + try: + if ( + upper_name == "RELATIONSHIP" + and "relationship_type" in table.column_names + ): + values = table.column("relationship_type").to_pylist() + distinct = sorted({str(value) for value in values if value is not None}) + if distinct: + if len(distinct) > max_value_samples: + display = ", ".join(distinct[:max_value_samples]) + ", ..." + else: + display = ", ".join(distinct) + extras.append(f"relationship_type values: {display}") + if upper_name == "RELATIONSHIP" and type_hints.get("relationship_types"): + allowed = ", ".join(type_hints["relationship_types"]) + extras.append(f"allowed relationship_type values: {allowed}") + except Exception: + pass + + lines.extend(f" {extra}" for extra in extras) + if not lines: + return "(no datasets available)" + return "\n".join(lines) + + +def build_type_hint_lines(type_hints: Mapping[str, tuple[str, ...]]) -> list[str]: + hints: list[str] = [] + entity_types = type_hints.get("entity_types") or () + relationship_types = type_hints.get("relationship_types") or () + if entity_types: + hints.append(f"entity_type values: {', '.join(entity_types)}") + if relationship_types: + hints.append(f"relationship_type values: {', '.join(relationship_types)}") + return hints + + +def _select_example_relationship_type( + type_hints: Mapping[str, tuple[str, ...]], +) -> str: + relationship_types = type_hints.get("relationship_types") or () + if relationship_types: + return relationship_types[0] + return "RELATIONSHIP_TYPE" + + +def build_query_prompt( + question: str, + schema_summary: str, + type_hint_lines: list[str], + type_hints: Mapping[str, tuple[str, ...]], + seed_entities: Sequence[Mapping[str, Any]] | None = None, + seed_neighbors: Sequence[Mapping[str, Any]] | None = None, +) -> str: + example_rel_type = _select_example_relationship_type(type_hints) + instruction_lines = [ + "You translate questions into Cypher for Lance graph datasets.", + ("Use the schema summary to craft queries that directly answer the question."), + ( + " • Use the schema summary and allowed relationship_type values to " + "identify candidate relationship directions and types." + ), + ( + " • When the schema lists relationship_type values and the question " + "does not narrow them down, treat the list as exhaustive and include " + "every value in your filter using OR clauses or " + "WHERE rel.relationship_type IN [...]." + ), + ( + "Always specify node labels and relationship types in MATCH patterns " + "that introduce aliases." + ), + "Supported constructs include:", + (" • MATCH (e:Entity) to scan entity rows (name, name_lower, entity_id)."), + ( + " • MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) to traverse " + "relationships (relationship_type column); `src` aligns with " + "`source_entity_id` and `dst` with `target_entity_id`." + ), + ( + " • Decide which node should be `src` versus `dst` based on the " + "relationship meaning in the question and schema hints." + ), + ( + " • Map natural language roles (team, person, product, etc.) to the " + "`entity_type` column so queries filter to the expected entities." + ), + " • Use WHERE e.column = 'value' for node-level filters.", + ( + " • Filter relationships with WHERE rel.relationship_type = 'VALUE' " + "or by comparing rel.source_entity_id / rel.target_entity_id; when the " + "question does not name a specific relationship type, include every " + "relevant value from the schema summary using OR clauses or " + "WHERE rel.relationship_type IN [...], explicitly note which values " + "you considered, and avoid emitting only a single guessed type." + ), + ( + " • Select columns using the aliases you define, such as e.name or " + "rel.relationship_type." + ), + ( + " • Avoid inventing relationship datasets; match RELATIONSHIP and " + "filter rel.relationship_type instead of [:TYPE]." + ), + ( + "Example: MATCH (part:Entity)-[rel:RELATIONSHIP]->(whole:Entity) " + f"WHERE rel.relationship_type = '{example_rel_type}' " + "RETURN part.name, whole.name." + ), + ( + "Example: MATCH (a:Entity)-[rel:RELATIONSHIP]->(b:Entity) WHERE " + "rel.relationship_type = 'TYPE_A' OR rel.relationship_type = 'TYPE_B' " + "RETURN a.name, b.name." + ), + ( + "Example: MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) WHERE " + "rel.relationship_type IN ['TYPE_A', 'TYPE_B', 'TYPE_C'] " + "RETURN src.name, dst.name." + ), + ( + "Example: MATCH (dst:Entity) WHERE dst.name_lower = 'acme corp' " + "RETURN dst.name, dst.entity_id." + ), + ( + f"Do not use relationship patterns like [:{example_rel_type}]; rely on " + "rel.relationship_type filters instead." + ), + ( + "Always emit at least one query when relevant data exists; only " + "return [] when it is impossible to answer." + ), + "Return ONLY a JSON array where each item has `cypher` and `description`.", + ] + if seed_entities: + instruction_lines.append( + "Prefer queries that start from the provided seed entities by referencing " + "their entity_id values before exploring related nodes." + ) + if seed_neighbors: + instruction_lines.extend( + [ + ( + "Use the provided seed neighbor relationships to decide " + "relationship direction." + ), + ( + " • Each neighbor entry includes a `direction` field: 'outgoing' " + "means the seed entity is the relationship source; 'incoming' " + "means the seed entity is the target." + ), + ( + " • Build MATCH patterns accordingly, e.g., outgoing -> " + "(seed)-[rel:RELATIONSHIP]->(neighbor); incoming -> " + "(neighbor)-[rel:RELATIONSHIP]->(seed)." + ), + ] + ) + instructions = "\n".join(instruction_lines) + + if type_hint_lines: + hint_block = "\n".join(f" • {line}" for line in type_hint_lines) + instructions = "\n".join( + [ + instructions, + "Allowed labels and type values:", + hint_block, + ] + ) + + prompt_parts = [ + instructions, + f"Schema summary:\n{schema_summary}", + ] + if seed_entities: + seed_lines = [] + for item in seed_entities: + similarity = item.get("similarity") + if isinstance(similarity, (int, float)): + score = f"{similarity:.3f}" + else: + score = "n/a" + display_name = str(item.get("name") or "(unknown)") + seed_lines.append( + ( + f"- {display_name} " + f"(entity_id={item.get('entity_id')}, " + f"entity_type={item.get('entity_type')}, similarity={score})" + ) + ) + prompt_parts.append( + "Seed entities discovered via embedding similarity:\n" + + "\n".join(seed_lines) + ) + if seed_neighbors: + neighbor_lines: list[str] = [] + for entry in seed_neighbors: + direction = str(entry.get("direction") or "outgoing") + seed_name = str( + entry.get("seed_name") or entry.get("seed_entity_id") or "(seed)" + ) + neighbor_name = str( + entry.get("neighbor_name") + or entry.get("neighbor_entity_id") + or "(neighbor)" + ) + rel_type = entry.get("relationship_type") or "RELATIONSHIP" + description = entry.get("relationship_description") or "" + seed_id = entry.get("seed_entity_id") + neighbor_id = entry.get("neighbor_entity_id") + if direction.lower() == "incoming": + arrow = f"{neighbor_name} -[{rel_type}]-> {seed_name}" + else: + arrow = f"{seed_name} -[{rel_type}]-> {neighbor_name}" + line = ( + f"- {arrow} (seed_entity_id={seed_id}, " + f"neighbor_entity_id={neighbor_id}, direction={direction}" + ) + if description: + line += f", description={description}" + line += ")" + neighbor_lines.append(line) + prompt_parts.append( + "Seed neighbor relationships (match patterns to respect direction):\n" + + "\n".join(neighbor_lines) + ) + prompt_parts.extend( + [ + f"Question:\n{question}", + "JSON:", + ] + ) + return "\n\n".join(prompt_parts) + + +def build_answer_prompt( + question: str, + schema_summary: str, + query_results: list[dict[str, Any]], +) -> str: + sections = [ + "You are a graph analysis assistant.", + "Provide a concise answer using the query results.", + "If the data is insufficient, state that clearly.", + "Schema summary:", + schema_summary, + "Query results:", + ] + for idx, item in enumerate(query_results, 1): + sections.append(f"Query {idx}: {item['cypher']}") + if item.get("description"): + sections.append(f"Description: {item['description']}") + if "error" in item: + sections.append(f"Error: {item['error']}") + else: + import json as _json + + rows_json = _json.dumps(item.get("rows", []), ensure_ascii=False, indent=2) + sections.append(f"Rows: {rows_json}") + if item.get("truncated"): + sections.append("(results truncated)") + sections.append(f"Question: {question}") + sections.append("Answer:") + return "\n".join(sections) diff --git a/python/python/knowledge_graph/llm/qa.py b/python/python/knowledge_graph/llm/qa.py new file mode 100644 index 000000000..68602ccc9 --- /dev/null +++ b/python/python/knowledge_graph/llm/qa.py @@ -0,0 +1,301 @@ +"""QA orchestration: planning, execution, and answer synthesis.""" + +from __future__ import annotations + +import json +import logging +import re +from typing import TYPE_CHECKING, Any + +from .. import extraction as kg_extraction +from .llm_utils import create_llm_client, load_llm_options, resolve_embedding_generator +from .prompts import ( + build_answer_prompt, + build_query_prompt, + build_type_hint_lines, + summarize_schema, +) +from .semantic import collect_seed_neighbors, find_seed_entities + +if TYPE_CHECKING: + from ..service import LanceKnowledgeGraph + from ..types import PlanStep, QueryResult, SeedEntity, SeedNeighbor + +LOGGER = logging.getLogger(__name__) + +DEFAULT_SEED_COUNT = 5 +DEFAULT_SEED_NEIGHBOR_LIMIT = 50 + + +def ask_question( + question: str, + service: LanceKnowledgeGraph, + *, + llm_model: str, + llm_temperature: float, + llm_config_path, + embedding_model: str | None, +) -> str: + client_options = load_llm_options(llm_config_path) + llm_client = create_llm_client( + llm_model=llm_model, + llm_temperature=llm_temperature, + llm_options=client_options, + ) + embedding_generator = resolve_embedding_generator( + model_name=embedding_model, options=client_options + ) + + seed_limit = DEFAULT_SEED_COUNT + + seed_entities: list[SeedEntity] = find_seed_entities( + question, + service, + embedding_generator, + limit=seed_limit, + ) + seed_neighbors: list[SeedNeighbor] = collect_seed_neighbors( + service, + seed_entities, + limit=DEFAULT_SEED_NEIGHBOR_LIMIT, + ) + + schema_summary = summarize_schema(service) + type_hints = service.store.config.type_hints() + allowed_relationship_types = tuple( + str(t) for t in (type_hints.get("relationship_types") or ()) + ) + if not allowed_relationship_types: + discovered = _discover_relationship_types(service) + allowed_relationship_types = tuple(discovered) + if discovered: + LOGGER.debug( + "Discovered relationship_type values from dataset: %s", + ", ".join(discovered), + ) + type_hint_lines = build_type_hint_lines(type_hints) + query_prompt = build_query_prompt( + question, + schema_summary, + type_hint_lines, + type_hints, + seed_entities, + seed_neighbors, + ) + + raw_plan = llm_client.complete(query_prompt) + plan_payload = kg_extraction.parse_llm_json(raw_plan) + query_plan: list[PlanStep] = extract_query_plan(plan_payload) + if allowed_relationship_types and query_plan: + query_plan = _normalize_relationship_types( + query_plan, allowed_relationship_types + ) + # Log normalized cypher for debugging + for step in query_plan: + LOGGER.debug("Normalized Cypher: %s", step.get("cypher")) + + if not query_plan and not seed_entities: + return "Unable to generate Cypher queries for the question." + + execution_results: list[QueryResult] = [] + if seed_entities: + execution_results.append( + { + "cypher": "(semantic search)", + "description": ( + "Top seed entities retrieved via embedding similarity search." + ), + "rows": seed_entities, + "truncated": False, + } + ) + if seed_neighbors: + execution_results.append( + { + "cypher": "(seed expansion)", + "description": ("Neighboring entities connected to the seed entities."), + "rows": seed_neighbors, + "truncated": bool( + DEFAULT_SEED_NEIGHBOR_LIMIT + and len(seed_neighbors) >= DEFAULT_SEED_NEIGHBOR_LIMIT + ), + } + ) + if query_plan: + execution_results.extend(execute_queries(service, query_plan)) + + if not execution_results: + return "Unable to gather context for the question." + + answer_prompt = build_answer_prompt(question, schema_summary, execution_results) + raw_answer = llm_client.complete(answer_prompt) + return raw_answer.strip() + + +def extract_query_plan(payload: Any) -> list[PlanStep]: + items: list[Any] + if isinstance(payload, list): + items = payload + elif isinstance(payload, dict): + items = payload.get("queries") or payload.get("plan") or [] + else: + return [] + + plan: list[PlanStep] = [] + for entry in items: + if not isinstance(entry, dict): + continue + cypher = entry.get("cypher") or entry.get("query") + if not cypher: + continue + plan.append( + { + "cypher": cypher, + "description": entry.get("description", ""), + } + ) + return plan + + +def execute_queries( + service: LanceKnowledgeGraph, plan: list[PlanStep], max_rows: int = 20 +) -> list[QueryResult]: + results: list[QueryResult] = [] + for step in plan: + cypher = step["cypher"] + description = step.get("description", "") + entry: QueryResult = {"cypher": cypher, "description": description} + try: + table = service.run(cypher) + rows = table.to_pylist() if hasattr(table, "to_pylist") else [] + truncated = False + if isinstance(rows, list) and len(rows) > max_rows: + truncated = True + rows = rows[:max_rows] + entry["rows"] = rows + entry["truncated"] = truncated + preview = json.dumps(rows, ensure_ascii=False, indent=2) + if truncated: + logging.debug( + "Cypher result (truncated to %s rows): %s", + max_rows, + preview, + ) + else: + logging.debug("Cypher result rows: %s", preview) + logging.debug( + "Cypher execution", + extra={"lance_graph": {"cypher": cypher, "rows": rows}}, + ) + except Exception as exc: # pragma: no cover - runtime safety + entry["error"] = str(exc) + data_preview = {} + for name, table in service.load_tables(service.dataset_names()).items(): + if hasattr(table, "schema"): + schema_names = list(table.schema.names) + else: + schema_names = [] + try: + row_limit = min(max_rows, getattr(table, "num_rows", 0)) + sample_rows = table.slice(0, row_limit).to_pylist() + except Exception: + sample_rows = [] + data_preview[name] = { + "schema": schema_names, + "rows_preview": sample_rows, + } + dataset_summary = json.dumps(data_preview, ensure_ascii=False, indent=2) + logging.debug( + "Cypher execution error\nCypher: %s\nError: %s\nDatasets: %s", + cypher, + str(exc), + dataset_summary, + ) + results.append(entry) + return results + + +def _normalize_relationship_types( + plan: list[PlanStep], allowed_types: tuple[str, ...] +) -> list[PlanStep]: + """Normalize or correct relationship_type filters to allowed canonical values. + + - Case-insensitive normalization for known values + - If unknown literal is used, replace predicate with IN [allowed_types] + - Supports patterns with and without the `rel.` alias + - Cleans IN lists to allowed values; if none remain, falls back to full allowed list + """ + if not allowed_types: + return plan + lowered = {t.lower(): t for t in allowed_types} + + def normalize_literal(value: str) -> tuple[str, bool]: + k = value.strip().lower() + if k in lowered: + return lowered[k], True + return value, False + + # equality patterns + eq_pats = [ + re.compile(r"(rel\.relationship_type\s*=\s*)'([^']+)'", flags=re.IGNORECASE), + re.compile(r"(relationship_type\s*=\s*)'([^']+)'", flags=re.IGNORECASE), + ] + + def replace_eq(match: re.Match[str]) -> str: + prefix, literal = match.group(1), match.group(2) + normalized, known = normalize_literal(literal) + if known: + return f"{prefix}'{normalized}'" + # Unknown literal: broaden to IN of all allowed types + allowed_list = ", ".join(f"'{t}'" for t in allowed_types) + # build a neutral predicate using the same LHS + lhs = prefix[:-1].strip() + return f"{lhs} IN [{allowed_list}]" + + # IN list pattern, with and without alias + in_pats = [ + re.compile( + r"(rel\.relationship_type\s+IN\s*\[)([^\]]+)(\])", + flags=re.IGNORECASE, + ), + re.compile(r"(relationship_type\s+IN\s*\[)([^\]]+)(\])", flags=re.IGNORECASE), + ] + + def replace_in(match: re.Match[str]) -> str: + start, inner, end = match.group(1), match.group(2), match.group(3) + items = re.findall(r"'([^']*)'", inner) + normalized_items: list[str] = [] + for it in items: + norm, known = normalize_literal(it) + if known: + normalized_items.append(f"'{norm}'") + if not normalized_items: + normalized_items = [f"'{t}'" for t in allowed_types] + return f"{start}{', '.join(normalized_items)}{end}" + + normalized_plan: list[PlanStep] = [] + for step in plan: + cypher = step.get("cypher", "") + for pat in eq_pats: + cypher = pat.sub(replace_eq, cypher) + for pat in in_pats: + cypher = pat.sub(replace_in, cypher) + normalized_step: PlanStep = { + "cypher": cypher, + "description": step.get("description", ""), + } + normalized_plan.append(normalized_step) + return normalized_plan + + +def _discover_relationship_types(service: LanceKnowledgeGraph) -> list[str]: + """Discover distinct relationship_type values from the dataset as a fallback.""" + try: + table = service.load_table("RELATIONSHIP") + if "relationship_type" in table.column_names: + values = table.column("relationship_type").to_pylist() + distinct = sorted({str(v) for v in values if v is not None and str(v)}) + return distinct + except Exception as exc: # pragma: no cover - defensive + LOGGER.debug("Unable to discover relationship types: %s", exc) + return [] diff --git a/python/python/knowledge_graph/llm/semantic.py b/python/python/knowledge_graph/llm/semantic.py new file mode 100644 index 000000000..e1accdce7 --- /dev/null +++ b/python/python/knowledge_graph/llm/semantic.py @@ -0,0 +1,151 @@ +"""Semantic helpers: seed entity search and neighbor collection.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Mapping, Sequence + +from ..embeddings import EmbeddingGenerator, cosine_similarity + +if TYPE_CHECKING: + from ..service import LanceKnowledgeGraph + from ..types import SeedEntity, SeedNeighbor + +LOGGER = logging.getLogger(__name__) + + +def find_seed_entities( + question: str, + service: "LanceKnowledgeGraph", + embedding_generator: EmbeddingGenerator | None, + *, + limit: int, +) -> list["SeedEntity"]: + if not embedding_generator: + return [] + prepared_question = question.strip() + if not prepared_question: + return [] + if limit <= 0: + return [] + if not service.has_dataset("Entity"): + return [] + try: + question_vector = embedding_generator.embed_one(prepared_question) + except Exception as exc: # pragma: no cover - defensive logging path + LOGGER.warning("Failed to embed question for semantic search: %s", exc) + return [] + if not question_vector: + return [] + try: + question_vector = [float(value) for value in question_vector] + except (TypeError, ValueError): + LOGGER.warning("Question embedding returned non-numeric values.") + return [] + try: + entity_table = service.load_table("Entity") + except Exception as exc: + LOGGER.warning("Unable to load Entity dataset for semantic search: %s", exc) + return [] + seeds: list[SeedEntity] = [] + for row in entity_table.to_pylist(): + embedding = row.get("embedding") + if not isinstance(embedding, (list, tuple)): + continue + try: + vector = [float(value) for value in embedding] + except (TypeError, ValueError): + continue + try: + similarity = float(cosine_similarity(question_vector, vector)) + except Exception: + similarity = 0.0 + entity_id = row.get("entity_id") + if not entity_id: + continue + seeds.append( + { + "entity_id": entity_id, + "name": row.get("name"), + "entity_type": row.get("entity_type"), + "similarity": similarity, + "context": row.get("context"), + } + ) + seeds.sort(key=lambda item: item.get("similarity", 0.0), reverse=True) + if limit and len(seeds) > limit: + seeds = seeds[:limit] + return seeds + + +def collect_seed_neighbors( + service: "LanceKnowledgeGraph", + seed_entities: Sequence["SeedEntity"], + *, + limit: int, +) -> list["SeedNeighbor"]: + if not seed_entities: + return [] + if not (service.has_dataset("Entity") and service.has_dataset("RELATIONSHIP")): + return [] + try: + entity_rows = service.load_table("Entity").to_pylist() + relationship_rows = service.load_table("RELATIONSHIP").to_pylist() + except Exception as exc: + LOGGER.warning("Unable to load datasets for neighbor expansion: %s", exc) + return [] + + id_to_entity: dict[str, Mapping[str, Any]] = {} + for entity in entity_rows: + entity_id = entity.get("entity_id") + if entity_id: + id_to_entity[str(entity_id)] = entity + + seed_ids = { + str(seed.get("entity_id")) for seed in seed_entities if seed.get("entity_id") + } + if not seed_ids: + return [] + + neighbors: list[SeedNeighbor] = [] + for relation in relationship_rows: + source_id = relation.get("source_entity_id") + target_id = relation.get("target_entity_id") + if source_id in seed_ids or target_id in seed_ids: + if source_id in seed_ids: + direction = "outgoing" + seed_id = str(source_id) + neighbor_id = str(target_id) if target_id else "" + else: + direction = "incoming" + seed_id = str(target_id) + neighbor_id = str(source_id) if source_id else "" + if not neighbor_id: + continue + seed_entity = id_to_entity.get(seed_id, {}) + neighbor_entity = id_to_entity.get(neighbor_id, {}) + neighbors.append( + { + "seed_entity_id": seed_id, + "seed_name": seed_entity.get("name"), + "seed_entity_type": seed_entity.get("entity_type"), + "neighbor_entity_id": neighbor_id, + "neighbor_name": neighbor_entity.get("name"), + "neighbor_entity_type": neighbor_entity.get("entity_type"), + "relationship_type": relation.get("relationship_type"), + "relationship_description": relation.get("description"), + "direction": direction, + } + ) + if not neighbors: + return [] + neighbors.sort( + key=lambda item: ( + str(item.get("seed_name") or ""), + str(item.get("neighbor_name") or ""), + str(item.get("relationship_type") or ""), + ) + ) + if limit and len(neighbors) > limit: + return neighbors[:limit] + return neighbors diff --git a/python/python/knowledge_graph/main.py b/python/python/knowledge_graph/main.py index 570fadd2a..b2586d03c 100644 --- a/python/python/knowledge_graph/main.py +++ b/python/python/knowledge_graph/main.py @@ -2,40 +2,30 @@ from __future__ import annotations -import argparse -import hashlib -import json import logging import sys -from dataclasses import asdict, is_dataclass -from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional, Sequence - -import yaml +from typing import TYPE_CHECKING, Optional, Sequence from . import extraction as kg_extraction -from .config import KnowledgeGraphConfig -from .embeddings import ( - DEFAULT_EMBEDDING_MODEL, - EmbeddingGenerator, - cosine_similarity, -) -from .service import LanceKnowledgeGraph +from .cli.ingest import extract_and_add, preview_extraction +from .cli.interactive import _execute_query, list_datasets, run_interactive +from .cli.parser import build_parser +from .cli.runtime import configure_logging as cli_configure_logging +from .cli.runtime import load_config as cli_load_config +from .cli.runtime import load_service as cli_load_service from .store import LanceGraphStore if TYPE_CHECKING: - import pyarrow as pa + import argparse - from .extractors import ExtractionResult + from .config import KnowledgeGraphConfig + from .service import LanceKnowledgeGraph LOGGER = logging.getLogger(__name__) -DEFAULT_SEED_COUNT = 5 -DEFAULT_SEED_NEIGHBOR_LIMIT = 50 - -def init_graph(config: KnowledgeGraphConfig) -> None: +def init_graph(config: "KnowledgeGraphConfig") -> None: """Initialize the on-disk storage and scaffold the schema file.""" config.ensure_directories() LanceGraphStore(config).ensure_layout() @@ -70,927 +60,37 @@ def init_graph(config: KnowledgeGraphConfig) -> None: print(f"Created schema template at {schema_path}") -def list_datasets(config: KnowledgeGraphConfig) -> None: - """List the Lance datasets available under the configured root.""" - store = LanceGraphStore(config) - store.ensure_layout() - datasets = store.list_datasets() - if not datasets: - print("No Lance datasets found. Load data or run extraction first.") - return - print("Available Lance datasets:") - for name, path in sorted(datasets.items()): - print(f" - {name}: {path}") - - -def run_interactive(service: LanceKnowledgeGraph) -> None: - """Enter an interactive shell for issuing Cypher queries.""" - print("Lance Knowledge Graph interactive shell") - print("Type ':help' for commands, or 'quit' to exit.") - - while True: - try: - text = input("kg> ").strip() - except EOFError: - print() - break - - if not text: - continue - lowered = text.lower() - if lowered in {"quit", "exit", "q"}: - break - if text.startswith(":"): - _handle_command(text, service) - continue - - _execute_query(service, text) - - -def _handle_command(command: str, service: LanceKnowledgeGraph) -> None: - """Handle meta-commands in the interactive shell.""" - cmd = command.strip() - if cmd in {":help", ":h"}: - print("Commands:") - print(" :help Show this message") - print(" :datasets List persisted Lance datasets") - print(" :config Show the configured node/relationship mappings") - print(" quit/exit/q Leave the shell") - return - if cmd in {":datasets", ":ls"}: - list_datasets(service.store.config) - return - if cmd in {":config", ":schema"}: - _print_config_summary(service) - return - print(f"Unknown command: {command}") - - -def _print_config_summary(service: LanceKnowledgeGraph) -> None: - """Print a brief summary of the graph configuration.""" - config = service.config - # GraphConfig does not currently expose direct iterators; rely on repr. - print("Graph configuration:") - print(f" {config!r}") - - -def _execute_query(service: LanceKnowledgeGraph, statement: str) -> None: - """Execute a single Cypher statement and print results.""" - try: - result = service.run(statement) - except Exception as exc: # pragma: no cover - CLI feedback path - print(f"Query failed: {exc}", file=sys.stderr) - return - - _print_table(result) - - -def _print_table(table: "pa.Table") -> None: - """Render a PyArrow table in a simple textual format.""" - if table.num_rows == 0: - print("(no rows)") - return - - column_names = table.column_names - columns = [table.column(i).to_pylist() for i in range(len(column_names))] - widths = [] - for name, values in zip(column_names, columns): - str_values = ["" if value is None else str(value) for value in values] - if str_values: - width = max(len(name), *(len(value) for value in str_values)) - else: - width = len(name) - widths.append(width) - - header = " | ".join(name.ljust(width) for name, width in zip(column_names, widths)) - separator = "-+-".join("-" * width for width in widths) - print(header) - print(separator) - for row_values in zip(*columns): - str_row = ["" if value is None else str(value) for value in row_values] - line = " | ".join(value.ljust(width) for value, width in zip(str_row, widths)) - print(line) - - -def preview_extraction(source: str, extractor: kg_extraction.BaseExtractor) -> None: - """Preview extracted knowledge from a text source or inline text.""" - text = _resolve_text_input(source) - result = kg_extraction.preview_extraction(text, extractor=extractor) - print(json.dumps(_result_to_dict(result), indent=2)) - - -def extract_and_add( - source: str, - service: LanceKnowledgeGraph, - extractor: kg_extraction.BaseExtractor, - *, - embedding_generator: EmbeddingGenerator | None = None, -) -> None: - """Extract knowledge and append it to the backing graph.""" - import pyarrow as pa - - text = _resolve_text_input(source) - result = kg_extraction.preview_extraction(text, extractor=extractor) - entity_rows, name_to_id = _prepare_entity_rows( - result.entities, embedding_generator=embedding_generator - ) - relationships = result.relationships - - if not entity_rows and not relationships: - print("No candidate entities or relationships detected.") - return - - if entity_rows: - entity_table = pa.Table.from_pylist(entity_rows) - service.upsert_table("Entity", entity_table, merge=True) - message = f"Upserted {entity_table.num_rows} entity rows into dataset 'Entity'." - print(message) - - relationship_rows = _prepare_relationship_rows( - relationships, - name_to_id, - embedding_generator=embedding_generator, - ) - if relationship_rows: - rel_table = pa.Table.from_pylist(relationship_rows) - service.upsert_table("RELATIONSHIP", rel_table, merge=True) - message = ( - "Upserted " - f"{rel_table.num_rows} relationship rows into dataset " - "'RELATIONSHIP'." - ) - print(message) - - def ask_question( question: str, - service: LanceKnowledgeGraph, - args: argparse.Namespace, + service: "LanceKnowledgeGraph", + args: "argparse.Namespace", ) -> None: """Answer a natural-language question using the graph via LLM-assisted Cypher.""" - client_options = _load_llm_options(args.llm_config) - llm_client = _create_llm_client(args, options=client_options) - embedding_generator = _resolve_embedding_generator(args, options=client_options) - seed_limit = getattr(args, "seed_count", DEFAULT_SEED_COUNT) - try: - seed_limit = int(seed_limit) - except (TypeError, ValueError): - seed_limit = DEFAULT_SEED_COUNT - if seed_limit < 0: - seed_limit = 0 + from .llm.qa import ask_question as qa_ask - seed_entities = _find_seed_entities( + answer = qa_ask( question, service, - embedding_generator, - limit=seed_limit, - ) - seed_neighbors = _collect_seed_neighbors( - service, - seed_entities, - limit=DEFAULT_SEED_NEIGHBOR_LIMIT, - ) - - schema_summary = _summarize_schema(service) - type_hints = service.store.config.type_hints() - type_hint_lines = _build_type_hint_lines(type_hints) - query_prompt = _build_query_prompt( - question, - schema_summary, - type_hint_lines, - type_hints, - seed_entities, - seed_neighbors, - ) - - raw_plan = llm_client.complete(query_prompt) - plan_payload = kg_extraction.parse_llm_json(raw_plan) - query_plan = _extract_query_plan(plan_payload) - - if not query_plan and not seed_entities: - print("Unable to generate Cypher queries for the question.") - return - - execution_results: list[dict[str, Any]] = [] - if seed_entities: - execution_results.append( - { - "cypher": "(semantic search)", - "description": ( - "Top seed entities retrieved via embedding similarity search." - ), - "rows": seed_entities, - "truncated": False, - } - ) - if seed_neighbors: - execution_results.append( - { - "cypher": "(seed expansion)", - "description": ("Neighboring entities connected to the seed entities."), - "rows": seed_neighbors, - "truncated": bool( - DEFAULT_SEED_NEIGHBOR_LIMIT - and len(seed_neighbors) >= DEFAULT_SEED_NEIGHBOR_LIMIT - ), - } - ) - if query_plan: - execution_results.extend(_execute_queries(service, query_plan)) - - if not execution_results: - print("Unable to gather context for the question.") - return - - answer_prompt = _build_answer_prompt(question, schema_summary, execution_results) - raw_answer = llm_client.complete(answer_prompt) - - print(raw_answer.strip()) - - -def _resolve_text_input(raw: str) -> str: - """Load text from a file if it exists, otherwise treat the string as content.""" - candidate = Path(raw) - if candidate.exists(): - if candidate.is_dir(): - raise IsADirectoryError(f"Expected text file, got directory: {candidate}") - return candidate.read_text(encoding="utf-8") - return raw - - -def _ensure_dict(item: object) -> dict: - if is_dataclass(item): - return asdict(item) # type: ignore[arg-type] - if isinstance(item, dict): - return item - raise TypeError(f"Unsupported extraction item type: {type(item)!r}") - - -def _result_to_dict(result: "ExtractionResult") -> dict[str, list[dict]]: - return { - "entities": [asdict(entity) for entity in result.entities], - "relationships": [asdict(rel) for rel in result.relationships], - } - - -def _prepare_entity_rows( - entities: list[Any], - *, - embedding_generator: EmbeddingGenerator | None = None, -) -> tuple[list[dict[str, Any]], dict[str, str]]: - rows: list[dict[str, Any]] = [] - name_to_id: dict[str, str] = {} - for entity in entities: - payload = _ensure_dict(entity) - name = str(payload.get("name", "")).strip() - entity_type = str( - payload.get("entity_type") or payload.get("type") or "" - ).strip() - if not name: - continue - base = f"{name}|{entity_type}".encode("utf-8") - entity_id = hashlib.md5(base).hexdigest() - payload["entity_id"] = entity_id - payload["entity_type"] = entity_type or "UNKNOWN" - payload["name_lower"] = name.lower() - rows.append(payload) - name_to_id.setdefault(name.lower(), entity_id) - if embedding_generator and rows: - _assign_embeddings( - rows, - embedding_generator, - _format_entity_embedding_input, - ) - return rows, name_to_id - - -def _prepare_relationship_rows( - relationships: list[Any], - name_to_id: dict[str, str], - *, - embedding_generator: EmbeddingGenerator | None = None, -) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - for relation in relationships: - payload = _ensure_dict(relation) - source_name = str( - payload.get("source_entity_name") or payload.get("source") or "" - ).strip() - target_name = str( - payload.get("target_entity_name") or payload.get("target") or "" - ).strip() - source_id = name_to_id.get(source_name.lower()) - target_id = name_to_id.get(target_name.lower()) - if not (source_id and target_id): - continue - payload["source_entity_id"] = source_id - payload["target_entity_id"] = target_id - payload["relationship_type"] = ( - payload.get("relationship_type") or payload.get("type") or "RELATED_TO" - ) - payload.setdefault("source_entity_name", source_name) - payload.setdefault("target_entity_name", target_name) - rows.append(payload) - if embedding_generator and rows: - _assign_embeddings( - rows, - embedding_generator, - _format_relationship_embedding_input, - ) - return rows - - -def _assign_embeddings( - rows: list[dict[str, Any]], - embedding_generator: EmbeddingGenerator, - formatter: Callable[[Mapping[str, Any]], str], -) -> None: - texts: list[str] = [] - indices: list[int] = [] - for idx, row in enumerate(rows): - text = formatter(row) - if text: - texts.append(text) - indices.append(idx) - if not texts: - return - try: - vectors = embedding_generator.embed(texts) - except Exception as exc: # pragma: no cover - defensive logging path - LOGGER.warning("Failed to generate embeddings: %s", exc) - return - if len(vectors) != len(indices): - LOGGER.warning( - "Mismatch between embedding count and row count: expected %s, got %s", - len(indices), - len(vectors), - ) - return - for idx, vector in zip(indices, vectors): - rows[idx]["embedding"] = vector - - -def _format_entity_embedding_input(row: Mapping[str, Any]) -> str: - name = str(row.get("name", "")).strip() - entity_type = str(row.get("entity_type", "")).strip() - context = str(row.get("context", "")).strip() - pieces = [] - if name: - pieces.append(name) - if entity_type: - pieces.append(f"Type: {entity_type}") - if context: - pieces.append(f"Context: {context}") - return " | ".join(pieces) - - -def _format_relationship_embedding_input(row: Mapping[str, Any]) -> str: - source = str(row.get("source_entity_name") or row.get("source") or "").strip() - target = str(row.get("target_entity_name") or row.get("target") or "").strip() - relationship_type = str(row.get("relationship_type", "")).strip() - description = str(row.get("description", "")).strip() - core: list[str] = [] - if source or target: - if relationship_type: - core.append(f"{source} -[{relationship_type}]-> {target}".strip()) - else: - core.append(f"{source} -> {target}".strip()) - if description: - core.append(f"Description: {description}") - return " | ".join(part for part in core if part) - - -def _find_seed_entities( - question: str, - service: LanceKnowledgeGraph, - embedding_generator: EmbeddingGenerator | None, - *, - limit: int = DEFAULT_SEED_COUNT, -) -> list[dict[str, Any]]: - if not embedding_generator: - return [] - prepared_question = question.strip() - if not prepared_question: - return [] - if limit <= 0: - return [] - if not service.has_dataset("Entity"): - return [] - try: - question_vector = embedding_generator.embed_one(prepared_question) - except Exception as exc: # pragma: no cover - defensive logging path - LOGGER.warning("Failed to embed question for semantic search: %s", exc) - return [] - if not question_vector: - return [] - try: - question_vector = [float(value) for value in question_vector] - except (TypeError, ValueError): - LOGGER.warning("Question embedding returned non-numeric values.") - return [] - try: - entity_table = service.load_table("Entity") - except Exception as exc: - LOGGER.warning("Unable to load Entity dataset for semantic search: %s", exc) - return [] - seeds: list[dict[str, Any]] = [] - for row in entity_table.to_pylist(): - embedding = row.get("embedding") - if not isinstance(embedding, (list, tuple)): - continue - try: - vector = [float(value) for value in embedding] - except (TypeError, ValueError): - continue - try: - similarity = float(cosine_similarity(question_vector, vector)) - except Exception: - similarity = 0.0 - entity_id = row.get("entity_id") - if not entity_id: - continue - seeds.append( - { - "entity_id": entity_id, - "name": row.get("name"), - "entity_type": row.get("entity_type"), - "similarity": similarity, - "context": row.get("context"), - } - ) - seeds.sort(key=lambda item: item.get("similarity", 0.0), reverse=True) - if limit and len(seeds) > limit: - seeds = seeds[:limit] - return seeds - - -def _collect_seed_neighbors( - service: LanceKnowledgeGraph, - seed_entities: Sequence[Mapping[str, Any]], - *, - limit: int = DEFAULT_SEED_NEIGHBOR_LIMIT, -) -> list[dict[str, Any]]: - if not seed_entities: - return [] - if not (service.has_dataset("Entity") and service.has_dataset("RELATIONSHIP")): - return [] - try: - entity_rows = service.load_table("Entity").to_pylist() - relationship_rows = service.load_table("RELATIONSHIP").to_pylist() - except Exception as exc: - LOGGER.warning("Unable to load datasets for neighbor expansion: %s", exc) - return [] - - id_to_entity: dict[str, Mapping[str, Any]] = {} - for entity in entity_rows: - entity_id = entity.get("entity_id") - if entity_id: - id_to_entity[str(entity_id)] = entity - - seed_ids = { - str(seed.get("entity_id")) for seed in seed_entities if seed.get("entity_id") - } - if not seed_ids: - return [] - - neighbors: list[dict[str, Any]] = [] - for relation in relationship_rows: - source_id = relation.get("source_entity_id") - target_id = relation.get("target_entity_id") - if source_id in seed_ids or target_id in seed_ids: - if source_id in seed_ids: - direction = "outgoing" - seed_id = str(source_id) - neighbor_id = str(target_id) if target_id else "" - else: - direction = "incoming" - seed_id = str(target_id) - neighbor_id = str(source_id) if source_id else "" - if not neighbor_id: - continue - seed_entity = id_to_entity.get(seed_id, {}) - neighbor_entity = id_to_entity.get(neighbor_id, {}) - neighbors.append( - { - "seed_entity_id": seed_id, - "seed_name": seed_entity.get("name"), - "seed_entity_type": seed_entity.get("entity_type"), - "neighbor_entity_id": neighbor_id, - "neighbor_name": neighbor_entity.get("name"), - "neighbor_entity_type": neighbor_entity.get("entity_type"), - "relationship_type": relation.get("relationship_type"), - "relationship_description": relation.get("description"), - "direction": direction, - } - ) - if not neighbors: - return [] - neighbors.sort( - key=lambda item: ( - str(item.get("seed_name") or ""), - str(item.get("neighbor_name") or ""), - str(item.get("relationship_type") or ""), - ) - ) - if limit and len(neighbors) > limit: - return neighbors[:limit] - return neighbors - - -def _summarize_schema( - service: LanceKnowledgeGraph, - max_columns: int = 20, - max_value_samples: int = 10, -) -> str: - type_hints = service.store.config.type_hints() - lines = [] - for name in service.dataset_names(): - try: - table = service.load_table(name) - except Exception: - continue - if hasattr(table, "schema"): - columns = list(getattr(table.schema, "names", [])) - else: - columns = [] - if len(columns) > max_columns: - columns = columns[:max_columns] + ["..."] - column_summary = f"- {name}: {', '.join(columns)}" - lines.append(column_summary) - - # Provide helpful value samples for well-known columns. - extras: list[str] = [] - upper_name = name.upper() - if upper_name == "ENTITY" and type_hints.get("entity_types"): - allowed = ", ".join(type_hints["entity_types"]) - extras.append(f"allowed entity_type values: {allowed}") - try: - if ( - upper_name == "RELATIONSHIP" - and "relationship_type" in table.column_names - ): - values = table.column("relationship_type").to_pylist() - distinct = sorted({str(value) for value in values if value is not None}) - if distinct: - if len(distinct) > max_value_samples: - display = ", ".join(distinct[:max_value_samples]) + ", ..." - else: - display = ", ".join(distinct) - extras.append(f"relationship_type values: {display}") - if upper_name == "RELATIONSHIP" and type_hints.get("relationship_types"): - allowed = ", ".join(type_hints["relationship_types"]) - extras.append(f"allowed relationship_type values: {allowed}") - except Exception: - pass - - lines.extend(f" {extra}" for extra in extras) - if not lines: - return "(no datasets available)" - return "\n".join(lines) - - -def _build_type_hint_lines(type_hints: Mapping[str, tuple[str, ...]]) -> list[str]: - hints: list[str] = [] - entity_types = type_hints.get("entity_types") or () - relationship_types = type_hints.get("relationship_types") or () - if entity_types: - hints.append(f"entity_type values: {', '.join(entity_types)}") - if relationship_types: - hints.append(f"relationship_type values: {', '.join(relationship_types)}") - return hints - - -def _select_example_relationship_type( - type_hints: Mapping[str, tuple[str, ...]], -) -> str: - relationship_types = type_hints.get("relationship_types") or () - if relationship_types: - return relationship_types[0] - return "RELATIONSHIP_TYPE" - - -def _build_query_prompt( - question: str, - schema_summary: str, - type_hint_lines: list[str], - type_hints: Mapping[str, tuple[str, ...]], - seed_entities: Sequence[Mapping[str, Any]] | None = None, - seed_neighbors: Sequence[Mapping[str, Any]] | None = None, -) -> str: - example_rel_type = _select_example_relationship_type(type_hints) - instruction_lines = [ - "You translate questions into Cypher for Lance graph datasets.", - ("Use the schema summary to craft queries that directly answer the question."), - ( - " • Use the schema summary and allowed relationship_type values to " - "identify candidate relationship directions and types." - ), - ( - " • When the schema lists relationship_type values and the question " - "does not narrow them down, treat the list as exhaustive and include " - "every value in your filter using OR clauses or " - "WHERE rel.relationship_type IN [...]." - ), - ( - "Always specify node labels and relationship types in MATCH patterns " - "that introduce aliases." - ), - "Supported constructs include:", - (" • MATCH (e:Entity) to scan entity rows (name, name_lower, entity_id)."), - ( - " • MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) to traverse " - "relationships (relationship_type column); `src` aligns with " - "`source_entity_id` and `dst` with `target_entity_id`." - ), - ( - " • Decide which node should be `src` versus `dst` based on the " - "relationship meaning in the question and schema hints." - ), - ( - " • Map natural language roles (team, person, product, etc.) to the " - "`entity_type` column so queries filter to the expected entities." - ), - " • Use WHERE e.column = 'value' for node-level filters.", - ( - " • Filter relationships with WHERE rel.relationship_type = 'VALUE' " - "or by comparing rel.source_entity_id / rel.target_entity_id; when the " - "question does not name a specific relationship type, include every " - "relevant value from the schema summary using OR clauses or " - "WHERE rel.relationship_type IN [...], explicitly note which values " - "you considered, and avoid emitting only a single guessed type." - ), - ( - " • Select columns using the aliases you define, such as e.name or " - "rel.relationship_type." - ), - ( - " • Avoid inventing relationship datasets; match RELATIONSHIP and " - "filter rel.relationship_type instead of [:TYPE]." - ), - ( - "Example: MATCH (part:Entity)-[rel:RELATIONSHIP]->(whole:Entity) " - f"WHERE rel.relationship_type = '{example_rel_type}' " - "RETURN part.name, whole.name." - ), - ( - "Example: MATCH (a:Entity)-[rel:RELATIONSHIP]->(b:Entity) WHERE " - "rel.relationship_type = 'TYPE_A' OR rel.relationship_type = 'TYPE_B' " - "RETURN a.name, b.name." - ), - ( - "Example: MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) WHERE " - "rel.relationship_type IN ['TYPE_A', 'TYPE_B', 'TYPE_C'] " - "RETURN src.name, dst.name." - ), - ( - "Example: MATCH (dst:Entity) WHERE dst.name_lower = 'acme corp' " - "RETURN dst.name, dst.entity_id." - ), - ( - f"Do not use relationship patterns like [:{example_rel_type}]; rely on " - "rel.relationship_type filters instead." - ), - ( - "Always emit at least one query when relevant data exists; only " - "return [] when it is impossible to answer." - ), - "Return ONLY a JSON array where each item has `cypher` and `description`.", - ] - if seed_entities: - instruction_lines.append( - "Prefer queries that start from the provided seed entities by referencing " - "their entity_id values before exploring related nodes." - ) - if seed_neighbors: - instruction_lines.extend( - [ - ( - "Use the provided seed neighbor relationships to decide " - "relationship direction." - ), - ( - " • Each neighbor entry includes a `direction` field: 'outgoing' " - "means the seed entity is the relationship source; 'incoming' " - "means the seed entity is the target." - ), - ( - " • Build MATCH patterns accordingly, e.g., outgoing -> " - "(seed)-[rel:RELATIONSHIP]->(neighbor); incoming -> " - "(neighbor)-[rel:RELATIONSHIP]->(seed)." - ), - ] - ) - instructions = "\n".join(instruction_lines) - - if type_hint_lines: - hint_block = "\n".join(f" • {line}" for line in type_hint_lines) - instructions = "\n".join( - [ - instructions, - "Allowed labels and type values:", - hint_block, - ] - ) - - prompt_parts = [ - instructions, - f"Schema summary:\n{schema_summary}", - ] - if seed_entities: - seed_lines = [] - for item in seed_entities: - similarity = item.get("similarity") - if isinstance(similarity, (int, float)): - score = f"{similarity:.3f}" - else: - score = "n/a" - display_name = str(item.get("name") or "(unknown)") - seed_lines.append( - ( - f"- {display_name} " - f"(entity_id={item.get('entity_id')}, " - f"entity_type={item.get('entity_type')}, similarity={score})" - ) - ) - prompt_parts.append( - "Seed entities discovered via embedding similarity:\n" - + "\n".join(seed_lines) - ) - if seed_neighbors: - neighbor_lines: list[str] = [] - for entry in seed_neighbors: - direction = str(entry.get("direction") or "outgoing") - seed_name = str( - entry.get("seed_name") or entry.get("seed_entity_id") or "(seed)" - ) - neighbor_name = str( - entry.get("neighbor_name") - or entry.get("neighbor_entity_id") - or "(neighbor)" - ) - rel_type = entry.get("relationship_type") or "RELATIONSHIP" - description = entry.get("relationship_description") or "" - seed_id = entry.get("seed_entity_id") - neighbor_id = entry.get("neighbor_entity_id") - if direction.lower() == "incoming": - arrow = f"{neighbor_name} -[{rel_type}]-> {seed_name}" - else: - arrow = f"{seed_name} -[{rel_type}]-> {neighbor_name}" - line = ( - f"- {arrow} (seed_entity_id={seed_id}, " - f"neighbor_entity_id={neighbor_id}, direction={direction}" - ) - if description: - line += f", description={description}" - line += ")" - neighbor_lines.append(line) - prompt_parts.append( - "Seed neighbor relationships (match patterns to respect direction):\n" - + "\n".join(neighbor_lines) - ) - prompt_parts.extend( - [ - f"Question:\n{question}", - "JSON:", - ] + llm_model=args.llm_model, + llm_temperature=args.llm_temperature, + llm_config_path=args.llm_config, + embedding_model=getattr(args, "embedding_model", None), ) - return "\n\n".join(prompt_parts) - - -def _extract_query_plan(payload: Any) -> list[dict[str, str]]: - items: list[Any] - if isinstance(payload, list): - items = payload - elif isinstance(payload, dict): - items = payload.get("queries") or payload.get("plan") or [] - else: - return [] - - plan: list[dict[str, str]] = [] - for entry in items: - if not isinstance(entry, dict): - continue - cypher = entry.get("cypher") or entry.get("query") - if not cypher: - continue - plan.append( - { - "cypher": cypher, - "description": entry.get("description", ""), - } - ) - return plan - + print(answer) -def _execute_queries( - service: LanceKnowledgeGraph, plan: list[dict[str, str]], max_rows: int = 20 -) -> list[dict[str, Any]]: - results: list[dict[str, Any]] = [] - for step in plan: - cypher = step["cypher"] - description = step.get("description", "") - entry: dict[str, Any] = {"cypher": cypher, "description": description} - try: - table = service.run(cypher) - rows = table.to_pylist() if hasattr(table, "to_pylist") else [] - truncated = False - if isinstance(rows, list) and len(rows) > max_rows: - truncated = True - rows = rows[:max_rows] - entry["rows"] = rows - entry["truncated"] = truncated - preview = json.dumps(rows, ensure_ascii=False, indent=2) - if truncated: - logging.debug( - "Cypher result (truncated to %s rows): %s", - max_rows, - preview, - ) - else: - logging.debug("Cypher result rows: %s", preview) - logging.debug( - "Cypher execution", - extra={"lance_graph": {"cypher": cypher, "rows": rows}}, - ) - except Exception as exc: # pragma: no cover - runtime safety - entry["error"] = str(exc) - data_preview = {} - for name, table in service.load_tables(service.dataset_names()).items(): - if hasattr(table, "schema"): - schema_names = list(table.schema.names) - else: - schema_names = [] - try: - row_limit = min(max_rows, getattr(table, "num_rows", 0)) - sample_rows = table.slice(0, row_limit).to_pylist() - except Exception: - sample_rows = [] - data_preview[name] = { - "schema": schema_names, - "rows_preview": sample_rows, - } - dataset_summary = json.dumps(data_preview, ensure_ascii=False, indent=2) - logging.debug( - "Cypher execution error\nCypher: %s\nError: %s\nDatasets: %s", - cypher, - str(exc), - dataset_summary, - ) - results.append(entry) - return results +def _load_config(args: "argparse.Namespace") -> "KnowledgeGraphConfig": + return cli_load_config(args) -def _build_answer_prompt( - question: str, - schema_summary: str, - query_results: list[dict[str, Any]], -) -> str: - sections = [ - "You are a graph analysis assistant.", - "Provide a concise answer using the query results.", - "If the data is insufficient, state that clearly.", - "Schema summary:", - schema_summary, - "Query results:", - ] - for idx, item in enumerate(query_results, 1): - sections.append(f"Query {idx}: {item['cypher']}") - if item.get("description"): - sections.append(f"Description: {item['description']}") - if "error" in item: - sections.append(f"Error: {item['error']}") - else: - rows_json = json.dumps(item.get("rows", []), ensure_ascii=False, indent=2) - sections.append(f"Rows: {rows_json}") - if item.get("truncated"): - sections.append("(results truncated)") - sections.append(f"Question: {question}") - sections.append("Answer:") - return "\n".join(sections) - - -def _load_config(args: argparse.Namespace) -> KnowledgeGraphConfig: - """Derive the configuration object from CLI arguments.""" - if args.root: - config = KnowledgeGraphConfig.from_root(Path(args.root)) - else: - config = KnowledgeGraphConfig.default() - if args.schema: - config = config.with_schema(Path(args.schema)) - return config +def _load_service(config: "KnowledgeGraphConfig"): + return cli_load_service(config) -def _load_service(config: KnowledgeGraphConfig) -> LanceKnowledgeGraph: - """Instantiate the knowledge graph service.""" - graph_config = config.load_graph_config() - storage = LanceGraphStore(config) - service = LanceKnowledgeGraph(graph_config, storage=storage) - service.ensure_initialized() - return service +def _resolve_extractor(args: "argparse.Namespace") -> kg_extraction.BaseExtractor: + from .llm.llm_utils import load_llm_options -def _resolve_extractor(args: argparse.Namespace) -> kg_extraction.BaseExtractor: - options = _load_llm_options(args.llm_config) + options = load_llm_options(args.llm_config) return kg_extraction.get_extractor( args.extractor, llm_model=args.llm_model, @@ -1000,152 +100,22 @@ def _resolve_extractor(args: argparse.Namespace) -> kg_extraction.BaseExtractor: def _resolve_embedding_generator( - args: argparse.Namespace, + args: "argparse.Namespace", *, - options: Optional[Mapping[str, Any]] = None, -) -> EmbeddingGenerator | None: - model = getattr(args, "embedding_model", None) - model_name = (model or "").strip() - if not model_name or model_name.lower() == "none": - return None - client_options = dict(options or _load_llm_options(args.llm_config)) - try: - return EmbeddingGenerator(model=model_name, client_options=client_options) - except RuntimeError as exc: - LOGGER.warning("Embeddings disabled: %s", exc) - return None + options: Optional[dict] = None, +): + from .llm.llm_utils import resolve_embedding_generator - -def _load_llm_options(path: Optional[Path]) -> dict: - if not path: - return {} - if not path.exists(): - raise FileNotFoundError(f"LLM config file not found: {path}") - with path.open("r", encoding="utf-8") as handle: - data = yaml.safe_load(handle) or {} - if not isinstance(data, dict): - raise ValueError("LLM config must be a mapping") - # Normalize nested http headers key to match OpenAI naming (default_headers) - if "default_headers" not in data and "http_headers" in data: - headers = data.pop("http_headers") - if isinstance(headers, dict): - data["default_headers"] = headers - return data - - -def _create_llm_client( - args: argparse.Namespace, - *, - options: Optional[Mapping[str, Any]] = None, -) -> kg_extraction.LLMClient: - resolved_options = dict(options or _load_llm_options(args.llm_config)) - return kg_extraction.get_llm_client( - llm_model=args.llm_model, - llm_temperature=args.llm_temperature, - llm_options=resolved_options, - ) + model = getattr(args, "embedding_model", None) + return resolve_embedding_generator(model_name=model, options=options) def _configure_logging(level: str) -> None: - normalized = level.upper() - numeric = getattr(logging, normalized, None) - if not isinstance(numeric, int): - raise ValueError(f"Invalid log level: {level}") - logging.basicConfig(level=numeric) + cli_configure_logging(level) -def _build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - prog="knowledge_graph", - description="Operate the Lance-backed knowledge graph.", - ) - parser.add_argument( - "--root", - type=Path, - help="Root directory for Lance datasets (default: ./knowledge_graph_data).", - ) - parser.add_argument( - "--schema", - type=Path, - help="Path to a YAML file describing node and relationship mappings.", - ) - parser.add_argument( - "--list-datasets", - action="store_true", - help="List stored Lance datasets and exit.", - ) - parser.add_argument( - "--extractor", - choices=["heuristic", "llm"], - default=kg_extraction.DEFAULT_STRATEGY, - help="Extraction strategy to use (default: llm).", - ) - parser.add_argument( - "--llm-model", - default="gpt-4o-mini", - help="LLM model identifier when using --extractor llm.", - ) - parser.add_argument( - "--llm-temperature", - type=float, - default=0.2, - help="Sampling temperature for --extractor llm (default: 0.2).", - ) - parser.add_argument( - "--llm-config", - type=Path, - help=( - "Optional YAML file with OpenAI client options (api_key, base_url, " - "headers, etc)." - ), - ) - parser.add_argument( - "--embedding-model", - default=DEFAULT_EMBEDDING_MODEL, - help=("OpenAI embedding model for semantic search (set to 'none' to disable)."), - ) - parser.add_argument( - "--seed-count", - type=int, - default=DEFAULT_SEED_COUNT, - help=( - "Maximum number of seed entities to surface from similarity search " - f"(default: {DEFAULT_SEED_COUNT})." - ), - ) - parser.add_argument( - "--log-level", - default="WARNING", - help="Logging level (DEBUG, INFO, WARNING, ERROR).", - ) - - group = parser.add_mutually_exclusive_group() - group.add_argument( - "--init", - action="store_true", - help="Initialize the knowledge graph storage.", - ) - group.add_argument( - "--extract-preview", - metavar="INPUT", - help="Preview extracted knowledge from a file path or raw text.", - ) - group.add_argument( - "--extract-and-add", - metavar="INPUT", - help="Extract knowledge from a file path or raw text and insert it.", - ) - group.add_argument( - "--ask", - metavar="QUESTION", - help="Ask a natural-language question over the knowledge graph.", - ) - parser.add_argument( - "query", - nargs="?", - help="Execute a single Cypher query against the Lance datasets.", - ) - return parser +def _build_parser() -> "argparse.ArgumentParser": + return build_parser() def main(argv: Optional[Sequence[str]] = None) -> int: diff --git a/python/python/knowledge_graph/semantic.py b/python/python/knowledge_graph/semantic.py new file mode 100644 index 000000000..58e03a404 --- /dev/null +++ b/python/python/knowledge_graph/semantic.py @@ -0,0 +1,151 @@ +"""Semantic helpers: seed entity search and neighbor collection.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Mapping, Sequence + +from .embeddings import EmbeddingGenerator, cosine_similarity + +if TYPE_CHECKING: + from .service import LanceKnowledgeGraph + from .types import SeedEntity, SeedNeighbor + +LOGGER = logging.getLogger(__name__) + + +def find_seed_entities( + question: str, + service: "LanceKnowledgeGraph", + embedding_generator: EmbeddingGenerator | None, + *, + limit: int, +) -> list["SeedEntity"]: + if not embedding_generator: + return [] + prepared_question = question.strip() + if not prepared_question: + return [] + if limit <= 0: + return [] + if not service.has_dataset("Entity"): + return [] + try: + question_vector = embedding_generator.embed_one(prepared_question) + except Exception as exc: # pragma: no cover - defensive logging path + LOGGER.warning("Failed to embed question for semantic search: %s", exc) + return [] + if not question_vector: + return [] + try: + question_vector = [float(value) for value in question_vector] + except (TypeError, ValueError): + LOGGER.warning("Question embedding returned non-numeric values.") + return [] + try: + entity_table = service.load_table("Entity") + except Exception as exc: + LOGGER.warning("Unable to load Entity dataset for semantic search: %s", exc) + return [] + seeds: list[SeedEntity] = [] + for row in entity_table.to_pylist(): + embedding = row.get("embedding") + if not isinstance(embedding, (list, tuple)): + continue + try: + vector = [float(value) for value in embedding] + except (TypeError, ValueError): + continue + try: + similarity = float(cosine_similarity(question_vector, vector)) + except Exception: + similarity = 0.0 + entity_id = row.get("entity_id") + if not entity_id: + continue + seeds.append( + { + "entity_id": entity_id, + "name": row.get("name"), + "entity_type": row.get("entity_type"), + "similarity": similarity, + "context": row.get("context"), + } + ) + seeds.sort(key=lambda item: item.get("similarity", 0.0), reverse=True) + if limit and len(seeds) > limit: + seeds = seeds[:limit] + return seeds + + +def collect_seed_neighbors( + service: "LanceKnowledgeGraph", + seed_entities: Sequence["SeedEntity"], + *, + limit: int, +) -> list["SeedNeighbor"]: + if not seed_entities: + return [] + if not (service.has_dataset("Entity") and service.has_dataset("RELATIONSHIP")): + return [] + try: + entity_rows = service.load_table("Entity").to_pylist() + relationship_rows = service.load_table("RELATIONSHIP").to_pylist() + except Exception as exc: + LOGGER.warning("Unable to load datasets for neighbor expansion: %s", exc) + return [] + + id_to_entity: dict[str, Mapping[str, Any]] = {} + for entity in entity_rows: + entity_id = entity.get("entity_id") + if entity_id: + id_to_entity[str(entity_id)] = entity + + seed_ids = { + str(seed.get("entity_id")) for seed in seed_entities if seed.get("entity_id") + } + if not seed_ids: + return [] + + neighbors: list[SeedNeighbor] = [] + for relation in relationship_rows: + source_id = relation.get("source_entity_id") + target_id = relation.get("target_entity_id") + if source_id in seed_ids or target_id in seed_ids: + if source_id in seed_ids: + direction = "outgoing" + seed_id = str(source_id) + neighbor_id = str(target_id) if target_id else "" + else: + direction = "incoming" + seed_id = str(target_id) + neighbor_id = str(source_id) if source_id else "" + if not neighbor_id: + continue + seed_entity = id_to_entity.get(seed_id, {}) + neighbor_entity = id_to_entity.get(neighbor_id, {}) + neighbors.append( + { + "seed_entity_id": seed_id, + "seed_name": seed_entity.get("name"), + "seed_entity_type": seed_entity.get("entity_type"), + "neighbor_entity_id": neighbor_id, + "neighbor_name": neighbor_entity.get("name"), + "neighbor_entity_type": neighbor_entity.get("entity_type"), + "relationship_type": relation.get("relationship_type"), + "relationship_description": relation.get("description"), + "direction": direction, + } + ) + if not neighbors: + return [] + neighbors.sort( + key=lambda item: ( + str(item.get("seed_name") or ""), + str(item.get("neighbor_name") or ""), + str(item.get("relationship_type") or ""), + ) + ) + if limit and len(neighbors) > limit: + return neighbors[:limit] + return neighbors diff --git a/python/python/knowledge_graph/types.py b/python/python/knowledge_graph/types.py new file mode 100644 index 000000000..392a463da --- /dev/null +++ b/python/python/knowledge_graph/types.py @@ -0,0 +1,38 @@ +"""Internal TypedDicts for knowledge graph CLI orchestration.""" + +from __future__ import annotations + +from typing import Any, NotRequired, TypedDict + + +class SeedEntity(TypedDict, total=False): + entity_id: str + name: NotRequired[str | None] + entity_type: NotRequired[str | None] + similarity: NotRequired[float] + context: NotRequired[str | None] + + +class SeedNeighbor(TypedDict, total=False): + seed_entity_id: str + seed_name: NotRequired[str | None] + seed_entity_type: NotRequired[str | None] + neighbor_entity_id: str + neighbor_name: NotRequired[str | None] + neighbor_entity_type: NotRequired[str | None] + relationship_type: NotRequired[str | None] + relationship_description: NotRequired[str | None] + direction: NotRequired[str] + + +class PlanStep(TypedDict, total=False): + cypher: str + description: NotRequired[str] + + +class QueryResult(TypedDict, total=False): + cypher: str + description: NotRequired[str] + rows: NotRequired[list[dict[str, Any]]] + truncated: NotRequired[bool] + error: NotRequired[str] diff --git a/python/uv.lock b/python/uv.lock index 1c0fd833f..f8f527be8 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -295,6 +295,8 @@ name = "lance-graph" source = { editable = "." } dependencies = [ { name = "fastapi" }, + { name = "lance" }, + { name = "openai" }, { name = "pyarrow" }, { name = "pydantic" }, { name = "pyyaml" }, @@ -306,12 +308,6 @@ dev = [ { name = "pyright" }, { name = "ruff" }, ] -lance-storage = [ - { name = "lance" }, -] -llm = [ - { name = "openai" }, -] tests = [ { name = "pandas" }, { name = "pyarrow" }, @@ -322,8 +318,8 @@ tests = [ [package.metadata] requires-dist = [ { name = "fastapi", specifier = ">=0.104.0" }, - { name = "lance", marker = "extra == 'lance-storage'", specifier = ">=0.17.0" }, - { name = "openai", marker = "extra == 'llm'", specifier = ">=1.52.0" }, + { name = "lance", specifier = ">=0.17.0" }, + { name = "openai", specifier = ">=1.52.0" }, { name = "pandas", marker = "extra == 'tests'" }, { name = "pyarrow", specifier = ">=14" }, { name = "pyarrow", marker = "extra == 'tests'", specifier = ">=14" }, @@ -335,7 +331,7 @@ requires-dist = [ { name = "ruff", marker = "extra == 'tests'" }, { name = "uvicorn", specifier = ">=0.24.0" }, ] -provides-extras = ["tests", "dev", "llm", "lance-storage"] +provides-extras = ["tests", "dev"] [[package]] name = "nodeenv"