mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-30 04:40:16 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fa8544c6d6 | ||
|
|
87649b7422 | ||
|
|
d91619f191 | ||
|
|
064a0db7e6 | ||
|
|
8214acc675 | ||
|
|
2bf55485ff | ||
|
|
1568237ce7 | ||
|
|
f6c9d50e03 | ||
|
|
d9117b7c2f | ||
|
|
0eabfb861e | ||
|
|
9f77dfb761 | ||
|
|
c990d09bd3 | ||
|
|
9ebacf43c3 | ||
|
|
7958ae78f6 | ||
|
|
2c61fe6cda | ||
|
|
92b850ac26 | ||
|
|
f7bd7016c5 | ||
|
|
8671385cbf | ||
|
|
b358acfabf | ||
|
|
a39ec5fd20 | ||
|
|
bbd6764215 | ||
|
|
1b0b0551db | ||
|
|
a6b102fa3d | ||
|
|
65d99f7f8a | ||
|
|
9b81137b26 | ||
|
|
653523efeb | ||
|
|
ba04421d9b | ||
|
|
5d3fe51dbd | ||
|
|
f20782f517 | ||
|
|
96dc5d754a | ||
|
|
cf84526cc7 | ||
|
|
5ad20abeab | ||
|
|
ade08a65ae | ||
|
|
fb25644fa7 | ||
|
|
63899f2427 | ||
|
|
fd6e058275 | ||
|
|
23d8207ef5 | ||
|
|
f2a11fc8ad | ||
|
|
c6316ba4bd | ||
|
|
b6d630fc74 | ||
|
|
3f2cb49e50 | ||
|
|
c7814616a9 | ||
|
|
531014fbda | ||
|
|
1cf9b34e3e | ||
|
|
2e81c86489 | ||
|
|
1690fec3f7 | ||
|
|
72a6ddb48f | ||
|
|
a5da533d55 | ||
|
|
be8856cfcf | ||
|
|
d2e599bcb0 | ||
|
|
05d0bbf86c | ||
|
|
dd7fcd3ddb | ||
|
|
43f55e1028 | ||
|
|
e20c522c62 | ||
|
|
fd9f0b2526 | ||
|
|
ccaadf6299 | ||
|
|
428fc3b83a | ||
|
|
09cf3ed132 | ||
|
|
58686d409b | ||
|
|
6d5fbc8b63 | ||
|
|
8c3f7f1f0a | ||
|
|
4acad23a4d | ||
|
|
cd1435ee10 | ||
|
|
68f0a1d4d9 | ||
|
|
a47274593b | ||
|
|
87a08e0240 | ||
|
|
1a2604255f | ||
|
|
94b312901b | ||
|
|
25fe95dd1a | ||
|
|
f338b66274 | ||
|
|
8b1cd47f51 | ||
|
|
48395b2f00 | ||
|
|
91ef2939c5 | ||
|
|
30d84c41ad | ||
|
|
a5c531fd29 | ||
|
|
976a20496d | ||
|
|
9bb94c2337 | ||
|
|
d8e04c29e9 |
@@ -61,6 +61,7 @@ wheels/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
.python-version
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
|
||||
@@ -5,6 +5,7 @@ repos:
|
||||
- id: trailing-whitespace
|
||||
- id: end-of-file-fixer
|
||||
- id: check-yaml
|
||||
exclude: 'neptune-setup\.yaml$'
|
||||
- id: check-json
|
||||
- id: check-toml
|
||||
- id: check-added-large-files
|
||||
@@ -49,9 +50,15 @@ repos:
|
||||
hooks:
|
||||
- id: yamllint
|
||||
args: ['-d', '{extends: default, rules: {line-length: {max: 120}}}']
|
||||
exclude: 'neptune-setup\.yaml$'
|
||||
|
||||
- repo: https://github.com/aws-cloudformation/cfn-lint
|
||||
rev: v1.43.3
|
||||
hooks:
|
||||
- id: cfn-lint
|
||||
files: 'neptune-setup\.yaml$'
|
||||
|
||||
# Removed slow hooks for faster development:
|
||||
# - mypy: Type checking (can be run manually or in CI)
|
||||
# - bandit: Security scanning (can be run separately)
|
||||
# - pytest: Testing (should be run manually, not on every commit)
|
||||
|
||||
|
||||
+172
-11
@@ -7,34 +7,195 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.2.3] - 2026-01-20
|
||||
|
||||
### Fixed
|
||||
- Fixed model switching bug in `TextEmbedder` where internal state was not cleared, preventing dynamic updates between `fastembed` and `sentence_transformers` (#160).
|
||||
- Implemented model-intrinsic embedding dimension detection in `TextEmbedder` to ensure consistency between models and vector databases.
|
||||
- Updated `set_model` to properly refresh configuration and dimensions during model switches.
|
||||
- **LLM Relation Extraction Parsing**:
|
||||
- Fixed relation extraction returning zero relations despite successful API calls to Groq and other providers
|
||||
- Normalized typed responses from instructor/OpenAI/Groq to consistent dict format before parsing
|
||||
- Added structured JSON fallback when typed generation yields zero relations to avoid silent empty outputs
|
||||
- Removed acceptance of extra kwargs (`max_tokens`, `max_entities_prompt`) from relation extraction internals
|
||||
- Filtered kwargs passed to provider LLM calls to only `temperature` and `verbose`
|
||||
- **API Parameter Handling**:
|
||||
- Limited kwargs forwarded in chunked extraction helper to prevent parameter leakage
|
||||
- Ensured minimal, safe parameters are passed to provider calls
|
||||
- **Pipeline Circular Import (Issues #192, #193)**:
|
||||
- Fixed circular import between `pipeline_builder` and `pipeline_validator` triggered during `semantica.pipeline` import
|
||||
- Lazy-loaded `PipelineValidator` inside `PipelineBuilder.__init__` and guarded type hints with `TYPE_CHECKING`
|
||||
- Ensured `from semantica.deduplication import DuplicateDetector` no longer fails even when pipeline module is imported
|
||||
- **JupyterLab Progress Output (Issue #181)**:
|
||||
- Added `SEMANTICA_DISABLE_JUPYTER_PROGRESS` environment variable to disable rich Jupyter/Colab progress tables
|
||||
- When enabled, progress falls back to console-style output, preventing infinite scrolling and JupyterLab out-of-memory errors
|
||||
|
||||
### Added
|
||||
- Added comprehensive unit test suite `tests/embeddings/test_model_switching.py` for verifying dynamic model transitions and dimension updates.
|
||||
- Fixed `TypeError: unhashable type: 'Entity'` in `GraphAnalyzer` when processing graphs with raw `Entity` objects or dictionaries in relationships (#159).
|
||||
- Robustified ID extraction across `CentralityCalculator`, `CommunityDetector`, and `ConnectivityAnalyzer` to handle various entity formats.
|
||||
- Improved `Entity` class hashability and equality logic in `utils/types.py`.
|
||||
- Added end-to-end integration test suite for Knowledge Graph pipeline validation (GraphBuilder -> EntityResolver -> GraphAnalyzer).
|
||||
- **Comprehensive Test Suite**:
|
||||
- - Added unit tests (`tests/test_relations_llm.py`) with mocked LLM provider covering both typed and structured response paths
|
||||
- - Added integration tests (`tests/integration/test_relations_groq.py`) for real Groq API calls with environment variable API key
|
||||
- - Tests validate relation extraction completion and result parsing across different response formats
|
||||
- **Amazon Neptune Dev Environment**:
|
||||
- - Added CloudFormation template (`cookbook/introduction/neptune-setup.yaml`) to provision a dev Neptune cluster with public endpoint and IAM auth enabled
|
||||
- - Documented deployment, cost estimates, and IAM User vs IAM Role best practices in `cookbook/introduction/21_Amazon_Neptune_Store.ipynb`
|
||||
- - Added `cfn-lint` to `.pre-commit-config.yaml` for validating CloudFormation templates while excluding `neptune-setup.yaml` from generic YAML linters
|
||||
- **Vector Store High-Performance Ingestion**:
|
||||
- - Added `VectorStore.add_documents` for high-throughput ingestion with automatic embedding generation, batching, and parallel processing
|
||||
- - Added `VectorStore.embed_batch` helper for generating embeddings for lists of texts without immediately storing them
|
||||
- - Enabled default parallel ingestion in `VectorStore` with `max_workers=6` for common workloads
|
||||
- - Added dedicated documentation page `docs/vector_store_usage.md` describing high-performance vector store usage and configuration
|
||||
- - Added `tests/vector_store/test_vector_store_parallel.py` covering parallel vs sequential performance, error handling, and edge cases for `add_documents` and `embed_batch`
|
||||
|
||||
### Changed
|
||||
- **Relation Extraction API**:
|
||||
- - Simplified parameter interface by removing unused kwargs that were previously ignored
|
||||
- - Improved error handling and verbose logging for debugging relation extraction issues
|
||||
- - Enhanced robustness of post-response parsing across different LLM providers
|
||||
- **Vector Store Defaults and Examples**:
|
||||
- - Standardized `VectorStore` default concurrency to `max_workers=6` for parallel ingestion
|
||||
- - Updated vector store reference documentation and usage guides to rely on implicit defaults instead of requiring manual `max_workers` configuration in examples
|
||||
|
||||
|
||||
## [0.2.2] - 2026-01-15
|
||||
|
||||
### Added
|
||||
- **Parallel Extraction Engine**:
|
||||
- Implemented high-throughput parallel batch processing across all core extractors (`NERExtractor`, `RelationExtractor`, `TripletExtractor`, `EventDetector`, `SemanticNetworkExtractor`) using `concurrent.futures.ThreadPoolExecutor`.
|
||||
- Added `max_workers` configuration parameter (default: 1) to all extractor `extract()` methods, allowing users to tune concurrency based on available CPU cores or API rate limits.
|
||||
- **Parallel Chunking**: Implemented parallel processing for large document chunking in `_extract_entities_chunked` and `_extract_relations_chunked`, significantly reducing latency for long-form text analysis.
|
||||
- **Thread-Safe Progress Tracking**: Enhanced `ProgressTracker` to handle concurrent updates from multiple threads without race conditions during batch processing.
|
||||
- **Semantic Extract Performance & Regression**:
|
||||
- Added edge-case regression suite covering max worker defaults, LLM prompt entity filtering, and extractor reuse.
|
||||
- Added a runnable real-use-case benchmark script for batch latency across `NERExtractor`, `RelationExtractor`, `TripletExtractor`, `EventDetector`, `SemanticAnalyzer`, and `SemanticNetworkExtractor`.
|
||||
- Added Groq LLM smoke tests that exercise LLM-based entities/relations/triplets when `GROQ_API_KEY` is available via environment configuration.
|
||||
|
||||
### Security
|
||||
- **Credential Sanitization**:
|
||||
- Removed hardcoded API keys from 8 cookbook notebooks to prevent secret leakage.
|
||||
- Enforced environment variable usage for `GROQ_API_KEY` across all examples.
|
||||
- **Secure Caching**:
|
||||
- Updated `ExtractionCache` to exclude sensitive parameters (e.g., `api_key`, `token`, `password`) from cache key generation, preventing secret leakage and enabling safe cache sharing.
|
||||
- Upgraded cache key hashing algorithm from MD5 to **SHA-256** for enhanced collision resistance and security.
|
||||
|
||||
### Changed
|
||||
- **Gemini SDK Migration**:
|
||||
- Migrated `GeminiProvider` to use the new `google-genai` SDK (v0.1.0+) to address deprecation warnings.
|
||||
- Implemented graceful fallback to `google.generativeai` for backward compatibility.
|
||||
- **Dependency Resolution**:
|
||||
- Pinned `opentelemetry-api` and `opentelemetry-sdk` to `1.37.0` to resolve pip conflicts.
|
||||
- Updated `protobuf` and `grpcio` constraints for better stability.
|
||||
- **Entity Filtering Scope**:
|
||||
- Removed entity filtering from non-LLM extraction flows to avoid accuracy regressions.
|
||||
- Applied entity downselection only to LLM relation prompt construction, while matching returned entities against the full original entity list.
|
||||
- **Batch Concurrency Defaults**:
|
||||
- Standardized `max_workers` defaulting across `semantic_extract` and tuned for low-latency: ML-backed methods default to single-worker, while pattern/regex/rules/LLM/huggingface methods use a higher parallelism default capped by CPU.
|
||||
- Raised the global `optimization.max_workers` default to 8 for better throughput on batch workloads.
|
||||
|
||||
### Performance
|
||||
- **Bottleneck Optimization (GitHub Issue #186)**:
|
||||
- **Resolved Bottleneck #1 (Sequential Processing)**: Replaced sequential `for` loops with parallel execution for both document-level batches and intra-document chunks.
|
||||
- **Performance Gains**: Achieved **~1.89x speedup** in real-world extraction scenarios (tested with Groq `llama-3.3-70b-versatile` on standard datasets).
|
||||
- **Initialization Optimization**: Refactored test suite to use class-level `setUpClass` for LLM provider initialization, eliminating redundant API client creation overhead.
|
||||
- **Low-Latency Entity Matching**:
|
||||
- Avoided heavyweight embedding stack imports on common matches by improving fast matching heuristics and short-circuiting before embedding similarity.
|
||||
- Optimized entity matching to prioritize exact/substring/word-boundary matches and only fall back to embedding similarity when needed, reducing CPU overhead in LLM relation/triplet mapping.
|
||||
|
||||
|
||||
## [0.2.1] - 2026-01-12
|
||||
|
||||
### Fixed
|
||||
- **LLM Output Stability (Bug #176)**:
|
||||
- Fixed incomplete JSON output issues by correctly propagating `max_tokens` parameter in `extract_relations_llm`.
|
||||
- Implemented automatic error handling that halves chunk sizes and retries when LLM context or output limits are exceeded.
|
||||
- Fixed `AttributeError` in provider integration by ensuring consistent parameter passing via `**kwargs`.
|
||||
- **Constraint Relaxations**:
|
||||
- Removed hardcoded `max_length` constraints from `Entity`, `Relation`, and `Triplet` classes to support long-form semantic extraction (e.g., long descriptions or names).
|
||||
- Fixed orchestrator lazy property initialization and configuration normalization logic in `Orchestrator`.
|
||||
- Resolved `AssertionError` in orchestrator tests by aligning test mocks with production component usage.
|
||||
- Fixed dependency compatibility issues by pinning `protobuf==4.25.3` and `grpcio==1.67.1`.
|
||||
- Fixed dependency compatibility issues by pinning `protobuf>=5.29.1,<7.0` and `grpcio>=1.71.2`.
|
||||
- Added missing dependencies `GitPython` and `chardet` to `pyproject.toml`.
|
||||
- Verified and aligned `FileObject.text` property usage in GraphRAG notebooks for consistent content decoding.
|
||||
|
||||
### Changed
|
||||
- **Chunking Defaults**:
|
||||
- Increased default `max_text_length` for auto-chunking to **64,000 characters** (from 32k/16k) for OpenAI, Anthropic, Gemini, Groq, and DeepSeek providers.
|
||||
- Unified chunking logic across `extract_entities_llm`, `extract_relations_llm`, and `extract_triplets_llm`.
|
||||
- **Groq Support**:
|
||||
- Standardized Groq provider defaults to use `llama-3.3-70b-versatile` with a 64k context window.
|
||||
- Added native support for `max_tokens` and `max_completion_tokens` to prevent output truncation.
|
||||
|
||||
### Added
|
||||
- **Testing**:
|
||||
- Added `tests/reproduce_issue_176.py` to validate `max_tokens` propagation and chunking behavior across all extractors.
|
||||
|
||||
|
||||
## [0.2.0] - 2026-01-10
|
||||
|
||||
### Added
|
||||
- **Amazon Neptune Support**:
|
||||
- Added `AmazonNeptuneStore` providing Amazon Neptune graph database integration via Bolt protocol and OpenCypher.
|
||||
- Implemented `NeptuneAuthTokenManager` extending Neo4j AuthManager for AWS IAM SigV4 signing with automatic token refresh.
|
||||
- Added robust connection handling: retry logic with backoff for transient errors (signature expired, connection closed) and driver recreation.
|
||||
- Added `graph-amazon-neptune` optional dependency group (boto3, neo4j).
|
||||
- Comprehensive test suite covering all GraphStore interface methods.
|
||||
- **Docling Integration**:
|
||||
- Added `DoclingParser` in `semantica.parse` for high-fidelity document parsing using the Docling library.
|
||||
- Supports multi-format parsing (PDF, DOCX, PPTX, XLSX, HTML, images) with superior table extraction and structure understanding.
|
||||
- Implemented as a standalone parser supporting local execution, OCR, and multiple export formats (Markdown, HTML, JSON).
|
||||
- **Robust Extraction Fallbacks**:
|
||||
- Implemented comprehensive fallback chains ("ML/LLM" -> "Pattern" -> "Last Resort") across `NERExtractor`, `RelationExtractor`, and `TripletExtractor` to prevent empty result lists.
|
||||
- Added "Last Resort" pattern matching in `NERExtractor` to identify capitalized words as generic entities when all other methods fail.
|
||||
- Added "Last Resort" adjacency-based relation extraction in `RelationExtractor` to create weak connections between adjacent entities if no relations are found.
|
||||
- Added fallback logic in `TripletExtractor` to convert relations to triplets or use rule-based extraction if standard methods fail.
|
||||
- **Provenance & Tracking**:
|
||||
- Added count tracking to batch processing logs in `NERExtractor`, `RelationExtractor`, and `TripletExtractor`.
|
||||
- Added `batch_index` and `document_id` to the metadata of all extracted entities, relations, triplets, semantic roles, and clusters for better traceability.
|
||||
- **Semantic Extract Improvements**:
|
||||
- Introduced `auto-chunking` for long text processing in LLM extraction methods (`extract_entities_llm`, `extract_relations_llm`, `extract_triplets_llm`).
|
||||
- Added `silent_fail` parameter to LLM extraction methods for configurable error handling.
|
||||
- Implemented robust JSON parsing and automatic retry logic (3 attempts with exponential backoff) in `BaseProvider` for all LLM providers.
|
||||
- Enhanced `GroqProvider` with better diagnostics and connectivity testing.
|
||||
- Added comprehensive entity, relation, and triplet deduplication for chunked extraction.
|
||||
- Added `semantica/semantic_extract/schemas.py` with canonical Pydantic models for consistent structured output.
|
||||
- **Testing**:
|
||||
- Added comprehensive robustness test suite `tests/semantic_extract/test_robustness_fallback.py` for validating extraction fallbacks and metadata propagation.
|
||||
- Added comprehensive unit test suite `tests/embeddings/test_model_switching.py` for verifying dynamic model transitions and dimension updates.
|
||||
- Added end-to-end integration test suite for Knowledge Graph pipeline validation (GraphBuilder -> EntityResolver -> GraphAnalyzer).
|
||||
- **Other**:
|
||||
- Added missing dependencies `GitPython` and `chardet` to `pyproject.toml`.
|
||||
- Robustified ID extraction across `CentralityCalculator`, `CommunityDetector`, and `ConnectivityAnalyzer` to handle various entity formats.
|
||||
- Improved `Entity` class hashability and equality logic in `utils/types.py`.
|
||||
|
||||
### Changed
|
||||
- **Deduplication & Conflict Logic**:
|
||||
- Removed internal deduplication logic from `NERExtractor`, `RelationExtractor`, and `TripletExtractor`.
|
||||
- Removed consistency/conflict checking from `ExtractionValidator` to defer to dedicated `semantica/conflicts` module.
|
||||
- Removed `_deduplicate_*` methods from `semantica/semantic_extract/methods.py`.
|
||||
- **Batch Processing & Consistency**:
|
||||
- Standardized batch processing across all extractors (`NERExtractor`, `RelationExtractor`, `TripletExtractor`, `SemanticNetworkExtractor`, `EventDetector`, `SemanticAnalyzer`, `CoreferenceResolver`) using a unified `extract`/`analyze`/`resolve` method pattern with progress tracking.
|
||||
- Added provenance metadata (`batch_index`, `document_id`) to `SemanticNetwork` nodes/edges, `Event` objects, `SemanticRole` results, `CoreferenceChain` mentions, and `SemanticCluster` (tracking source `document_ids`).
|
||||
- Updated `SemanticClusterer.cluster` and `SemanticAnalyzer.cluster_semantically` to accept list of dictionaries (with `content` and `id` keys) for better document tracking during clustering.
|
||||
- Removed legacy `check_triplet_consistency` from `TripletExtractor`.
|
||||
- Removed `validate_consistency` and `_check_consistency` from `ExtractionValidator`.
|
||||
- **Weighted Scoring**:
|
||||
- Clarified weighted confidence scoring (50% Method Confidence + 50% Type Similarity) in comments.
|
||||
- Explicitly labeled "Type Similarity" as "user-provided" in code comments to remove ambiguity.
|
||||
- **Refactoring**:
|
||||
- Fixed orchestrator lazy property initialization and configuration normalization logic in `Orchestrator`.
|
||||
- Verified and aligned `FileObject.text` property usage in GraphRAG notebooks for consistent content decoding.
|
||||
|
||||
### Fixed
|
||||
- Fixed a bug in `TripletExtractor` where the `validate_triplets` method was shadowed by an internal attribute.
|
||||
- Fixed incorrect `TextSplitter` import path in the `semantic_extract.methods` module.
|
||||
- **Critical Fixes**:
|
||||
- Resolved `NameError` in `extraction_validator.py` by adding missing `Union` import.
|
||||
- Resolved issues where extractors would return empty lists for valid input text when primary extraction methods failed.
|
||||
- Fixed metadata initialization issue in batch processing where `batch_index` and `document_id` were occasionally missing from extracted items.
|
||||
- Ensured `LLMExtraction` methods (`enhance_entities`, `enhance_relations`) return original input instead of failing or returning empty results when LLM providers are unavailable.
|
||||
- **Component Fixes**:
|
||||
- Fixed model switching bug in `TextEmbedder` where internal state was not cleared, preventing dynamic updates between `fastembed` and `sentence_transformers` (#160).
|
||||
- Implemented model-intrinsic embedding dimension detection in `TextEmbedder` to ensure consistency between models and vector databases.
|
||||
- Updated `set_model` to properly refresh configuration and dimensions during model switches.
|
||||
- Fixed `TypeError: unhashable type: 'Entity'` in `GraphAnalyzer` when processing graphs with raw `Entity` objects or dictionaries in relationships (#159).
|
||||
- Resolved `AssertionError` in orchestrator tests by aligning test mocks with production component usage.
|
||||
- Fixed dependency compatibility issues by pinning `protobuf==4.25.3` and `grpcio==1.67.1`.
|
||||
- Fixed a bug in `TripletExtractor` where the `validate_triplets` method was shadowed by an internal attribute.
|
||||
- Fixed incorrect `TextSplitter` import path in the `semantic_extract.methods` module.
|
||||
|
||||
## [0.1.1] - 2026-01-05
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
[](https://www.python.org/downloads/)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://pypi.org/project/semantica/)
|
||||
[](https://pypi.org/project/semantica/)
|
||||
[](https://pypi.org/project/semantica/)
|
||||
[](https://pepy.tech/project/semantica)
|
||||
[](https://discord.gg/pMHguUzG)
|
||||
@@ -28,7 +28,7 @@
|
||||
|
||||
*The missing fabric between raw data and AI engineering. A comprehensive open-source framework for building semantic layers and knowledge engineering systems that transform unstructured data into AI-ready knowledge — powering Knowledge Graph-Powered RAG (GraphRAG), AI Agents, Multi-Agent Systems, and AI applications with structured semantic knowledge.*
|
||||
|
||||
**100% Open Source** • **MIT Licensed** • **Latest Version: 0.1.1** • **Production Ready** • **Community Driven**
|
||||
**100% Open Source** • **MIT Licensed** • **Latest Version: 0.2.3** • **Production Ready** • **Community Driven**
|
||||
|
||||
[**Discord**](https://discord.gg/pMHguUzG)
|
||||
|
||||
@@ -360,7 +360,7 @@ results = vector_store.search(query="supply chain", top_k=5)
|
||||
|
||||
### Graph Store & Triplet Store
|
||||
|
||||
> **Neo4j, FalkorDB support** • **SPARQL queries** • **RDF triplets**
|
||||
> **Neo4j, FalkorDB, Amazon Neptune support** • **SPARQL queries** • **RDF triplets**
|
||||
|
||||
```python
|
||||
from semantica.graph_store import GraphStore
|
||||
@@ -370,6 +370,24 @@ from semantica.triplet_store import TripletStore
|
||||
graph_store = GraphStore(backend="neo4j", uri="bolt://localhost:7687", user="neo4j", password="password")
|
||||
graph_store.add_nodes([{"id": "n1", "labels": ["Person"], "properties": {"name": "Alice"}}])
|
||||
|
||||
# Amazon Neptune Graph Store (OpenCypher via HTTP with IAM Auth)
|
||||
neptune_store = GraphStore(
|
||||
backend="neptune",
|
||||
endpoint="your-cluster.us-east-1.neptune.amazonaws.com",
|
||||
port=8182,
|
||||
region="us-east-1",
|
||||
iam_auth=True, # Uses AWS credential chain (boto3, env vars, or IAM role)
|
||||
)
|
||||
|
||||
# Node Operations
|
||||
neptune_store.add_nodes([
|
||||
{"labels": ["Person"], "properties": {"id": "alice", "name": "Alice", "age": 30}},
|
||||
{"labels": ["Person"], "properties": {"id": "bob", "name": "Bob", "age": 25}},
|
||||
])
|
||||
|
||||
# Query Operations
|
||||
result = neptune_store.execute_query("MATCH (p:Person) RETURN p.name, p.age")
|
||||
|
||||
# Triplet Store (Blazegraph, Jena, RDF4J)
|
||||
triplet_store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999/blazegraph")
|
||||
triplet_store.add_triplet({"subject": "Alice", "predicate": "knows", "object": "Bob"})
|
||||
|
||||
+3
-3
@@ -26,10 +26,10 @@ Before releasing, ensure:
|
||||
|
||||
The project uses GitHub Actions for automated releases to PyPI.
|
||||
|
||||
1. **Tag the commit**: Create a new git tag for the version (e.g., `v0.1.1`).
|
||||
1.29. **Tag the commit**: Create a new git tag for the version (e.g., `v0.2.3`).
|
||||
```bash
|
||||
git tag -a v0.1.1 -m "Release v0.1.1"
|
||||
git push origin v0.1.1
|
||||
git tag -a v0.2.3 -m "Release v0.2.3"
|
||||
git push origin v0.2.3
|
||||
```
|
||||
2. **GitHub Action**: The `Release` workflow will automatically trigger, build the package, create a GitHub Release, and publish to PyPI using Trusted Publishing.
|
||||
|
||||
|
||||
@@ -6,6 +6,10 @@ We actively support the following versions of Semantica with security updates:
|
||||
|
||||
| Version | Supported |
|
||||
| ------- | ------------------ |
|
||||
| 0.2.3 | :white_check_mark: |
|
||||
| 0.2.2 | :white_check_mark: |
|
||||
| 0.2.1 | :white_check_mark: |
|
||||
| 0.2.0 | :white_check_mark: |
|
||||
| 0.1.1 | :white_check_mark: |
|
||||
| 0.1.0 | :white_check_mark: |
|
||||
| < 0.1.0 | :x: |
|
||||
|
||||
@@ -0,0 +1,725 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Amazon Neptune Graph Store\n",
|
||||
"\n",
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"This notebook covers the Amazon Neptune Database integration in Semantica. Amazon Neptune is a fully managed graph database service that supports both property graphs (via OpenCypher/Gremlin) and RDF graphs (via SPARQL).\n",
|
||||
"\n",
|
||||
"### Key Features\n",
|
||||
"\n",
|
||||
"- **IAM Authentication**: Secure access using AWS SigV4 signatures via AuthManager\n",
|
||||
"- **OpenCypher Support**: Query using standard OpenCypher syntax\n",
|
||||
"- **Bolt Protocol**: Uses Neo4j Bolt driver for efficient binary communication\n",
|
||||
"- **Native ~id Support**: Leverages Neptune's native element ID handling\n",
|
||||
"- **Full CRUD Operations**: Create, read, update, delete nodes and relationships\n",
|
||||
"- **Automatic Retry**: Built-in retry logic with exponential backoff for transient errors\n",
|
||||
"\n",
|
||||
"### Prerequisites\n",
|
||||
"\n",
|
||||
"- An Amazon Neptune Database cluster\n",
|
||||
"- AWS credentials configured (boto3, environment variables, or IAM role)\n",
|
||||
"- Network access to your Neptune cluster (VPC, security groups)\n",
|
||||
"\n",
|
||||
"#### Quick Setup with CloudFormation\n",
|
||||
"\n",
|
||||
"If you don't have a Neptune cluster, use the provided CloudFormation template to create one with a public endpoint and IAM authentication:\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"# Deploy the Neptune stack (takes ~15-20 minutes)\n",
|
||||
"aws cloudformation create-stack \\\n",
|
||||
" --stack-name semantica-neptune \\\n",
|
||||
" --template-body file://neptune-setup.yaml \\\n",
|
||||
" --capabilities CAPABILITY_NAMED_IAM\n",
|
||||
"\n",
|
||||
"# Wait for stack creation to complete\n",
|
||||
"aws cloudformation wait stack-create-complete --stack-name semantica-neptune\n",
|
||||
"\n",
|
||||
"# Get the outputs (endpoint, port, credentials)\n",
|
||||
"aws cloudformation describe-stacks --stack-name semantica-neptune \\\n",
|
||||
" --query 'Stacks[0].Outputs' --output table\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"The template creates:\n",
|
||||
"- VPC with public subnets and Internet Gateway\n",
|
||||
"- Neptune cluster (`db.t3.medium`) with IAM authentication enabled\n",
|
||||
"- IAM user with least-privilege access for OpenCypher queries\n",
|
||||
"- Security group allowing Bolt protocol (port 8182) access\n",
|
||||
"\n",
|
||||
"> ⚠️ **Security Note**: This template creates an IAM User with static access keys for simplicity in demo/test environments. For production use, we recommend IAM Roles (EC2 instance roles, ECS task roles, Lambda execution roles) which provide temporary credentials that are automatically rotated. The secret access key in the Cloudformation outputs is provided in plaintext to simplify initial setup - in production, use AWS Secrets Manager.\n",
|
||||
"\n",
|
||||
"**Outputs:**\n",
|
||||
"- `NeptuneEndpoint` - Cluster hostname (use as `NEPTUNE_ENDPOINT`)\n",
|
||||
"- `NeptunePort` - 8182 (use as `NEPTUNE_PORT`)\n",
|
||||
"- `AwsAccessKeyId` - IAM user access key (use as `AWS_ACCESS_KEY_ID`)\n",
|
||||
"- `AwsSecretAccessKey` - IAM user secret key in **plaintext** (use as `AWS_SECRET_ACCESS_KEY`)\n",
|
||||
"- `AwsRegion` - Deployment region (use as `AWS_REGION`)\n",
|
||||
"\n",
|
||||
"**Cleanup:**\n",
|
||||
"```bash\n",
|
||||
"aws cloudformation delete-stack --stack-name semantica-neptune\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"**Estimated Monthly Cost (approximately 100-105 USD/month at 100% utilization):**\n",
|
||||
"\n",
|
||||
"| Resource | Cost (USD) |\n",
|
||||
"| --- | --- |\n",
|
||||
"| Neptune db.t3.medium instance | ~96/month (0.132/hr) |\n",
|
||||
"| Storage (10 GB) | ~1/month |\n",
|
||||
"| I/O requests | ~1-5/month |\n",
|
||||
"| Public IPv4 address | ~3.60/month (0.005/hr) |\n",
|
||||
"| VPC, subnets, route tables, Internet Gateway, IAM | No Additional Charge |\n",
|
||||
"\n",
|
||||
"> **Free Tier**: New Neptune users get 30 days free (750 hours of db.t3.medium, 10M I/Os, 1 GB storage). Delete the stack when not in use to avoid charges.\n",
|
||||
"\n",
|
||||
"---"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Installation\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"# Install Semantica with Neptune support\n",
|
||||
"pip install semantica\n",
|
||||
"\n",
|
||||
"# Required dependencies (installed automatically)\n",
|
||||
"pip install boto3 neo4j\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install semantica"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Configuration\n",
|
||||
"\n",
|
||||
"Set your Neptune cluster endpoint and AWS credentials. Replace the placeholder values with your actual configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"# Neptune cluster configuration - REPLACE WITH YOUR VALUES\n",
|
||||
"# (Get these from CloudFormation stack outputs)\n",
|
||||
"os.environ[\"NEPTUNE_ENDPOINT\"] = \"your-cluster.us-east-1.neptune.amazonaws.com\"\n",
|
||||
"os.environ[\"NEPTUNE_PORT\"] = \"8182\"\n",
|
||||
"os.environ[\"AWS_REGION\"] = \"us-east-1\"\n",
|
||||
"\n",
|
||||
"# AWS credentials for IAM Authentication\n",
|
||||
"# Option 1: IAM User (static credentials from CloudFormation template)\n",
|
||||
"# os.environ[\"AWS_ACCESS_KEY_ID\"] = \"AKIA...\" # From AwsAccessKeyId output\n",
|
||||
"# os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"...\" # From AwsSecretAccessKey output\n",
|
||||
"# Note: No AWS_SESSION_TOKEN needed for IAM users\n",
|
||||
"\n",
|
||||
"# Option 2: IAM Role / Temporary credentials (e.g., STS AssumeRole, EC2 instance role)\n",
|
||||
"# os.environ[\"AWS_ACCESS_KEY_ID\"] = \"ASIA...\" # Temporary access key\n",
|
||||
"# os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"...\" # Temporary secret key\n",
|
||||
"# os.environ[\"AWS_SESSION_TOKEN\"] = \"...\" # REQUIRED for temporary credentials\n",
|
||||
"\n",
|
||||
"print(f\"Neptune Endpoint: {os.environ.get('NEPTUNE_ENDPOINT')}\")\n",
|
||||
"print(f\"AWS Region: {os.environ.get('AWS_REGION')}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 1: Initialize Neptune Store\n",
|
||||
"\n",
|
||||
"Initialize a connection to your Amazon Neptune cluster with IAM authentication."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from semantica.graph_store import GraphStore\n",
|
||||
"\n",
|
||||
"# Option 1: Using GraphStore factory (recommended)\n",
|
||||
"neptune_store = GraphStore(\n",
|
||||
" backend=\"neptune\",\n",
|
||||
" endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n",
|
||||
" port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n",
|
||||
" region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n",
|
||||
" iam_auth=True,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Connect to Neptune\n",
|
||||
"neptune_store.connect()\n",
|
||||
"print(\"Connected to Amazon Neptune!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Development/Testing Without IAM Auth\n",
|
||||
"\n",
|
||||
"For development or testing environments where IAM authentication is not required (e.g., Neptune notebooks or VPC-only access), you can disable IAM signing:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# For dev/test environments without IAM authentication\n",
|
||||
"neptune_store_dev = GraphStore(\n",
|
||||
" backend=\"neptune\",\n",
|
||||
" endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n",
|
||||
" port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n",
|
||||
" region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n",
|
||||
" iam_auth=False, # Disable IAM signing for dev/test\n",
|
||||
")\n",
|
||||
"neptune_store_dev.connect()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Authentication Options\n",
|
||||
"\n",
|
||||
"IAM Authentication (recommended for production) automatically uses the AWS credential chain:\n",
|
||||
"1. Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)\n",
|
||||
"2. AWS credentials file (~/.aws/credentials)\n",
|
||||
"3. IAM role (for EC2, Lambda, ECS)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 2: Node Operations\n",
|
||||
"\n",
|
||||
"### Creating Nodes\n",
|
||||
"\n",
|
||||
"Nodes represent entities in your graph. Each node can have:\n",
|
||||
"- **ID**: A unique identifier (custom or auto-generated UUID)\n",
|
||||
"- **Labels**: Categories/types (e.g., `Person`, `Company`)\n",
|
||||
"- **Properties**: Key-value pairs (e.g., `{\"name\": \"Alice\", \"age\": 30}`)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Create a single node with custom ID (id in properties)\n",
|
||||
"alice = neptune_store.create_node(\n",
|
||||
" labels=[\"Person\"],\n",
|
||||
" properties={\"id\": \"alice\", \"name\": \"Alice\", \"age\": 30, \"role\": \"Engineer\"}\n",
|
||||
")\n",
|
||||
"print(f\"Created node: {alice}\")\n",
|
||||
"\n",
|
||||
"# Create a node with auto-generated UUID (no id in properties)\n",
|
||||
"bob = neptune_store.create_node(\n",
|
||||
" labels=[\"Person\"],\n",
|
||||
" properties={\"name\": \"Bob\", \"age\": 25, \"role\": \"Designer\"}\n",
|
||||
")\n",
|
||||
"print(f\"Created node with UUID: {bob['id']}\")\n",
|
||||
"\n",
|
||||
"# Create a company node with auto-generated ID\n",
|
||||
"acme = neptune_store.create_node(\n",
|
||||
" labels=[\"Company\"],\n",
|
||||
" properties={\"name\": \"Acme Corp\", \"industry\": \"Technology\", \"founded\": 2010}\n",
|
||||
")\n",
|
||||
"print(f\"Created company: {acme}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Creating Multiple Nodes (Batch)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Batch create nodes for better performance\n",
|
||||
"# Include 'id' in properties for custom IDs\n",
|
||||
"nodes_data = [\n",
|
||||
" {\"labels\": [\"Person\"], \"properties\": {\"id\": \"charlie\", \"name\": \"Charlie\", \"age\": 35}},\n",
|
||||
" {\"labels\": [\"Person\"], \"properties\": {\"id\": \"diana\", \"name\": \"Diana\", \"age\": 28}},\n",
|
||||
" {\"labels\": [\"Location\"], \"properties\": {\"name\": \"San Francisco\", \"state\": \"CA\"}},\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"created_nodes = neptune_store.create_nodes(nodes_data)\n",
|
||||
"print(f\"Created {len(created_nodes)} nodes in batch\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Retrieving Nodes"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Get a specific node by ID\n",
|
||||
"alice_node = neptune_store.get_node(node_id=\"alice\")\n",
|
||||
"print(f\"Retrieved: {alice_node}\")\n",
|
||||
"\n",
|
||||
"# Get nodes by label\n",
|
||||
"people = neptune_store.get_nodes(labels=[\"Person\"], limit=10)\n",
|
||||
"print(f\"Found {len(people)} Person nodes:\")\n",
|
||||
"for person in people:\n",
|
||||
" print(f\" - {person.get('properties', {}).get('name')}\")\n",
|
||||
"\n",
|
||||
"# Get nodes by properties\n",
|
||||
"engineers = neptune_store.get_nodes(\n",
|
||||
" labels=[\"Person\"],\n",
|
||||
" properties={\"role\": \"Engineer\"},\n",
|
||||
" limit=5\n",
|
||||
")\n",
|
||||
"print(f\"Found {len(engineers)} engineers\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Updating Nodes"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Update node properties (merge mode - default)\n",
|
||||
"updated_alice = neptune_store.update_node(\n",
|
||||
" node_id=\"alice\",\n",
|
||||
" properties={\"age\": 31, \"department\": \"AI Research\"},\n",
|
||||
" merge=True\n",
|
||||
")\n",
|
||||
"print(f\"Updated Alice: {updated_alice}\")\n",
|
||||
"\n",
|
||||
"# Replace all properties (merge=False)\n",
|
||||
"# WARNING: This removes properties not in the update\n",
|
||||
"replaced = neptune_store.update_node(\n",
|
||||
" node_id=\"charlie\",\n",
|
||||
" properties={\"name\": \"Charlie\", \"age\": 36},\n",
|
||||
" merge=False\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Deleting Nodes"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Delete a node (with detach=True to also delete relationships)\n",
|
||||
"deleted = neptune_store.delete_node(node_id=\"diana\", detach=True)\n",
|
||||
"print(f\"Deleted diana: {deleted}\")\n",
|
||||
"\n",
|
||||
"# Without detach (fails if node has relationships)\n",
|
||||
"# neptune_store.delete_node(node_id=\"alice\", detach=False)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 3: Relationship Operations\n",
|
||||
"\n",
|
||||
"### Creating Relationships\n",
|
||||
"\n",
|
||||
"Relationships connect nodes and represent connections between entities."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Create a relationship between Alice and Acme\n",
|
||||
"works_at = neptune_store.create_relationship(\n",
|
||||
" start_node_id=\"alice\",\n",
|
||||
" end_node_id=acme[\"id\"],\n",
|
||||
" rel_type=\"WORKS_AT\",\n",
|
||||
" properties={\"since\": 2020, \"position\": \"Senior Engineer\"}\n",
|
||||
")\n",
|
||||
"print(f\"Created relationship: {works_at}\")\n",
|
||||
"\n",
|
||||
"# Create a KNOWS relationship between people\n",
|
||||
"knows_rel = neptune_store.create_relationship(\n",
|
||||
" start_node_id=\"alice\",\n",
|
||||
" end_node_id=bob[\"id\"],\n",
|
||||
" rel_type=\"KNOWS\",\n",
|
||||
" properties={\"since\": 2019}\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Retrieving Relationships"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Get all relationships for a node\n",
|
||||
"alice_rels = neptune_store.get_relationships(node_id=\"alice\", direction=\"both\")\n",
|
||||
"print(f\"Alice has {len(alice_rels)} relationships\")\n",
|
||||
"\n",
|
||||
"# Get outgoing relationships only\n",
|
||||
"outgoing = neptune_store.get_relationships(node_id=\"alice\", direction=\"out\")\n",
|
||||
"\n",
|
||||
"# Filter by relationship type\n",
|
||||
"works_rels = neptune_store.get_relationships(\n",
|
||||
" node_id=\"alice\",\n",
|
||||
" rel_type=\"WORKS_AT\",\n",
|
||||
" direction=\"out\"\n",
|
||||
")\n",
|
||||
"print(f\"Alice's work relationships: {len(works_rels)}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Deleting Relationships"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Delete a specific relationship by ID\n",
|
||||
"if works_at.get(\"id\"):\n",
|
||||
" deleted = neptune_store.delete_relationship(rel_id=works_at[\"id\"])\n",
|
||||
" print(f\"Deleted relationship: {deleted}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 4: OpenCypher Queries\n",
|
||||
"\n",
|
||||
"Amazon Neptune supports OpenCypher queries via the Bolt protocol. Execute complex graph patterns using standard Cypher syntax."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Simple query\n",
|
||||
"results = neptune_store.execute_query(\n",
|
||||
" \"MATCH (p:Person) RETURN p.name, p.age ORDER BY p.age\"\n",
|
||||
")\n",
|
||||
"print(\"People in the graph:\")\n",
|
||||
"for record in results.get(\"records\", []):\n",
|
||||
" print(f\" - {record.get('p.name')}: {record.get('p.age')} years old\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Using parameters (safer and more efficient)\n",
|
||||
"results = neptune_store.execute_query(\n",
|
||||
" \"MATCH (p:Person) WHERE p.age > $min_age RETURN p.name, p.age\",\n",
|
||||
" parameters={\"min_age\": 25}\n",
|
||||
")\n",
|
||||
"print(f\"People over 25: {len(results.get('records', []))}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Find relationships between nodes\n",
|
||||
"results = neptune_store.execute_query(\"\"\"\n",
|
||||
" MATCH (p:Person)-[r:WORKS_AT]->(c:Company)\n",
|
||||
" RETURN p.name as employee, c.name as company, r.since as start_year\n",
|
||||
"\"\"\")\n",
|
||||
"for record in results.get(\"records\", []):\n",
|
||||
" print(f\"{record['employee']} works at {record['company']} since {record['start_year']}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Count and aggregate\n",
|
||||
"results = neptune_store.execute_query(\"\"\"\n",
|
||||
" MATCH (p:Person)\n",
|
||||
" RETURN count(p) as total, avg(p.age) as avg_age, max(p.age) as max_age\n",
|
||||
"\"\"\")\n",
|
||||
"stats = results.get(\"records\", [{}])[0]\n",
|
||||
"print(f\"Total: {stats.get('total')}, Avg Age: {stats.get('avg_age'):.1f}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 5: Graph Analytics\n",
|
||||
"\n",
|
||||
"### Get Neighbors\n",
|
||||
"\n",
|
||||
"Traverse the graph to find connected nodes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Get immediate neighbors (depth=1)\n",
|
||||
"neighbors = neptune_store.get_neighbors(\n",
|
||||
" node_id=\"alice\",\n",
|
||||
" direction=\"both\",\n",
|
||||
" depth=1\n",
|
||||
")\n",
|
||||
"print(f\"Alice's direct neighbors: {len(neighbors)}\")\n",
|
||||
"\n",
|
||||
"# Get neighbors up to 2 hops away\n",
|
||||
"extended = neptune_store.get_neighbors(\n",
|
||||
" node_id=\"alice\",\n",
|
||||
" direction=\"out\",\n",
|
||||
" depth=2\n",
|
||||
")\n",
|
||||
"print(f\"Nodes within 2 hops: {len(extended)}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Shortest Path\n",
|
||||
"\n",
|
||||
"Find the shortest path between two nodes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Find shortest path\n",
|
||||
"path = neptune_store.shortest_path(\n",
|
||||
" start_node_id=\"alice\",\n",
|
||||
" end_node_id=\"charlie\",\n",
|
||||
" max_depth=5\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"if path:\n",
|
||||
" print(\"Path found!\")\n",
|
||||
" print(f\" Length: {path.get('length')}\")\n",
|
||||
" print(f\" Nodes: {len(path.get('nodes', []))}\")\n",
|
||||
" print(f\" Relationships: {len(path.get('relationships', []))}\")\n",
|
||||
"else:\n",
|
||||
" print(\"No path found between nodes\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 6: Graph Statistics\n",
|
||||
"\n",
|
||||
"Get comprehensive statistics about your graph."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Get graph statistics\n",
|
||||
"stats = neptune_store.get_stats()\n",
|
||||
"\n",
|
||||
"print(\"Graph Statistics:\")\n",
|
||||
"print(f\" Total nodes: {stats.get('node_count', 'N/A')}\")\n",
|
||||
"print(f\" Total relationships: {stats.get('relationship_count', 'N/A')}\")\n",
|
||||
"\n",
|
||||
"print(\"\\nNode labels:\")\n",
|
||||
"for label, count in stats.get('label_counts', {}).items():\n",
|
||||
" print(f\" - {label}: {count}\")\n",
|
||||
"\n",
|
||||
"print(\"\\nRelationship types:\")\n",
|
||||
"for rel_type, count in stats.get('relationship_type_counts', {}).items():\n",
|
||||
" print(f\" - {rel_type}: {count}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Step 7: Connection Management\n",
|
||||
"\n",
|
||||
"Always close connections when done to free resources."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Check connection status\n",
|
||||
"status = neptune_store.get_status()\n",
|
||||
"print(f\"Connection status: {status}\")\n",
|
||||
"\n",
|
||||
"# Close the connection\n",
|
||||
"neptune_store.close()\n",
|
||||
"print(\"Connection closed\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Neptune-Specific Considerations\n",
|
||||
"\n",
|
||||
"### Native Element IDs\n",
|
||||
"\n",
|
||||
"Neptune uses native `~id` for element identification. Include `id` in properties to set a custom ID:\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"# Create a node with custom ID (include 'id' in properties)\n",
|
||||
"node = neptune_store.create_node(\n",
|
||||
" labels=[\"Person\"],\n",
|
||||
" properties={\"id\": \"my-custom-id\", \"name\": \"Test\"}\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Create a node with auto-generated UUID (omit 'id' from properties)\n",
|
||||
"node = neptune_store.create_node(\n",
|
||||
" labels=[\"Person\"],\n",
|
||||
" properties={\"name\": \"Test\"}\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# The ID is used in id() function calls internally:\n",
|
||||
"# MATCH (n) WHERE id(n) = 'my-custom-id' RETURN n\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"### OpenCypher Considerations\n",
|
||||
"\n",
|
||||
"Amazon Neptune Database's OpenCypher implementation has some differences from Neo4j:\n",
|
||||
"\n",
|
||||
"1. **No `shortestPath()` function**: Use variable-length path patterns or `allShortestPaths()`\n",
|
||||
"2. **Labels syntax**: Use `labels(n)` function to retrieve node labels\n",
|
||||
"3. **Property updates**: Use `SET n += {props}` for merge behavior\n",
|
||||
"\n",
|
||||
"For the complete OpenCypher specification supported by Amazon Neptune Database, see the [AWS documentation](https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-opencypher.html).\n",
|
||||
"\n",
|
||||
"### Amazon Neptune Analytics\n",
|
||||
"\n",
|
||||
"For analytical (OLAP) workloads such as graph algorithms, aggregations, and large-scale traversals, consider [Amazon Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html). Neptune Analytics complements Neptune Database by providing optimized performance for analytical queries while Neptune Database is optimized for transactional (OLTP) workloads.\n",
|
||||
"\n",
|
||||
"### Performance Tips\n",
|
||||
"\n",
|
||||
"1. **Use batch operations** for creating multiple nodes/relationships\n",
|
||||
"2. **Use parameters** in queries to enable query caching\n",
|
||||
"3. **Limit result sets** with `LIMIT` clause"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Summary\n",
|
||||
"\n",
|
||||
"This notebook covered the Amazon Neptune Graph Store integration:\n",
|
||||
"\n",
|
||||
"- **IAM Authentication**: Secure AWS SigV4 signing\n",
|
||||
"- **CRUD Operations**: Full node and relationship management\n",
|
||||
"- **OpenCypher Queries**: Standard graph query language\n",
|
||||
"- **Graph Analytics**: Neighbors and shortest path algorithms\n",
|
||||
"- **Statistics & Monitoring**: Graph metrics and status\n",
|
||||
"\n",
|
||||
"### Key Takeaways\n",
|
||||
"\n",
|
||||
"- Neptune uses native `~id` for element identification\n",
|
||||
"- IAM authentication is recommended for production\n",
|
||||
"- Bolt protocol provides efficient binary query interface\n",
|
||||
"- Semantica abstracts Neptune-specific syntax differences\n",
|
||||
"\n",
|
||||
"### Next Steps\n",
|
||||
"\n",
|
||||
"- [Graph Store (Neo4j/FalkorDB)](09_Graph_Store.ipynb) - Compare with other backends\n",
|
||||
"- [Building Knowledge Graphs](07_Building_Knowledge_Graphs.ipynb) - Build production KGs\n",
|
||||
"- [Graph Analytics](10_Graph_Analytics.ipynb) - Advanced analytics algorithms"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python",
|
||||
"version": "3.9.0"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
AWSTemplateFormatVersion: '2010-09-09'
|
||||
Description: >
|
||||
Amazon Neptune cluster with public endpoint, IAM authentication, and least-privilege
|
||||
IAM user for Semantica cookbook. Uses db.t3.medium (most cost-effective Neptune instance type).
|
||||
|
||||
Parameters:
|
||||
EnvironmentName:
|
||||
Type: String
|
||||
Default: semantica-neptune
|
||||
Description: Environment name prefix for resource naming
|
||||
|
||||
Resources:
|
||||
# =============================================================================
|
||||
# VPC & NETWORKING
|
||||
# =============================================================================
|
||||
|
||||
VPC:
|
||||
Type: AWS::EC2::VPC
|
||||
Properties:
|
||||
CidrBlock: 10.0.0.0/16
|
||||
EnableDnsHostnames: true
|
||||
EnableDnsSupport: true
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-vpc
|
||||
|
||||
InternetGateway:
|
||||
Type: AWS::EC2::InternetGateway
|
||||
Properties:
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-igw
|
||||
|
||||
InternetGatewayAttachment:
|
||||
Type: AWS::EC2::VPCGatewayAttachment
|
||||
Properties:
|
||||
InternetGatewayId: !Ref InternetGateway
|
||||
VpcId: !Ref VPC
|
||||
|
||||
PublicSubnet1:
|
||||
Type: AWS::EC2::Subnet
|
||||
Properties:
|
||||
VpcId: !Ref VPC
|
||||
AvailabilityZone: !Select [0, !GetAZs '']
|
||||
CidrBlock: 10.0.1.0/24
|
||||
MapPublicIpOnLaunch: true
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-public-subnet-1
|
||||
|
||||
PublicSubnet2:
|
||||
Type: AWS::EC2::Subnet
|
||||
Properties:
|
||||
VpcId: !Ref VPC
|
||||
AvailabilityZone: !Select [1, !GetAZs '']
|
||||
CidrBlock: 10.0.2.0/24
|
||||
MapPublicIpOnLaunch: true
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-public-subnet-2
|
||||
|
||||
PublicRouteTable:
|
||||
Type: AWS::EC2::RouteTable
|
||||
Properties:
|
||||
VpcId: !Ref VPC
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-public-rt
|
||||
|
||||
DefaultPublicRoute:
|
||||
Type: AWS::EC2::Route
|
||||
DependsOn: InternetGatewayAttachment
|
||||
Properties:
|
||||
RouteTableId: !Ref PublicRouteTable
|
||||
DestinationCidrBlock: 0.0.0.0/0
|
||||
GatewayId: !Ref InternetGateway
|
||||
|
||||
PublicSubnet1RouteTableAssociation:
|
||||
Type: AWS::EC2::SubnetRouteTableAssociation
|
||||
Properties:
|
||||
RouteTableId: !Ref PublicRouteTable
|
||||
SubnetId: !Ref PublicSubnet1
|
||||
|
||||
PublicSubnet2RouteTableAssociation:
|
||||
Type: AWS::EC2::SubnetRouteTableAssociation
|
||||
Properties:
|
||||
RouteTableId: !Ref PublicRouteTable
|
||||
SubnetId: !Ref PublicSubnet2
|
||||
|
||||
# =============================================================================
|
||||
# SECURITY GROUP
|
||||
# =============================================================================
|
||||
|
||||
NeptuneSecurityGroup:
|
||||
Type: AWS::EC2::SecurityGroup
|
||||
Properties:
|
||||
GroupName: !Sub ${EnvironmentName}-neptune-sg
|
||||
GroupDescription: Security group for Neptune cluster - allows Bolt protocol access
|
||||
VpcId: !Ref VPC
|
||||
SecurityGroupIngress:
|
||||
- IpProtocol: tcp
|
||||
FromPort: 8182
|
||||
ToPort: 8182
|
||||
CidrIp: 0.0.0.0/0
|
||||
Description: Allow Bolt protocol access from anywhere
|
||||
SecurityGroupEgress:
|
||||
- IpProtocol: -1
|
||||
CidrIp: 0.0.0.0/0
|
||||
Description: Allow all outbound traffic
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-neptune-sg
|
||||
|
||||
# =============================================================================
|
||||
# NEPTUNE CLUSTER
|
||||
# =============================================================================
|
||||
|
||||
NeptuneSubnetGroup:
|
||||
Type: AWS::Neptune::DBSubnetGroup
|
||||
Properties:
|
||||
DBSubnetGroupDescription: Subnet group for Neptune cluster
|
||||
DBSubnetGroupName: !Sub ${EnvironmentName}-subnet-group
|
||||
SubnetIds:
|
||||
- !Ref PublicSubnet1
|
||||
- !Ref PublicSubnet2
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-subnet-group
|
||||
|
||||
NeptuneCluster:
|
||||
Type: AWS::Neptune::DBCluster
|
||||
Properties:
|
||||
DBClusterIdentifier: !Sub ${EnvironmentName}-cluster
|
||||
DBSubnetGroupName: !Ref NeptuneSubnetGroup
|
||||
VpcSecurityGroupIds:
|
||||
- !Ref NeptuneSecurityGroup
|
||||
EngineVersion: '1.4.6.3'
|
||||
IamAuthEnabled: true
|
||||
StorageEncrypted: true
|
||||
DeletionProtection: false
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-cluster
|
||||
|
||||
NeptuneInstance:
|
||||
Type: AWS::Neptune::DBInstance
|
||||
Properties:
|
||||
DBInstanceIdentifier: !Sub ${EnvironmentName}-instance
|
||||
DBInstanceClass: db.t3.medium
|
||||
DBClusterIdentifier: !Ref NeptuneCluster
|
||||
PubliclyAccessible: true
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-instance
|
||||
|
||||
# =============================================================================
|
||||
# IAM USER WITH LEAST PRIVILEGES
|
||||
# =============================================================================
|
||||
|
||||
NeptuneUser:
|
||||
Type: AWS::IAM::User
|
||||
Properties:
|
||||
UserName: !Sub ${EnvironmentName}-user
|
||||
Tags:
|
||||
- Key: Name
|
||||
Value: !Sub ${EnvironmentName}-user
|
||||
|
||||
NeptuneUserPolicy:
|
||||
Type: AWS::IAM::Policy
|
||||
Properties:
|
||||
PolicyName: !Sub ${EnvironmentName}-neptune-access
|
||||
Users:
|
||||
- !Ref NeptuneUser
|
||||
PolicyDocument:
|
||||
Version: '2012-10-17'
|
||||
Statement:
|
||||
- Sid: NeptuneDataAccess
|
||||
Effect: Allow
|
||||
Action:
|
||||
- neptune-db:connect
|
||||
- neptune-db:ReadDataViaQuery
|
||||
- neptune-db:WriteDataViaQuery
|
||||
- neptune-db:DeleteDataViaQuery
|
||||
Resource: !Sub
|
||||
- arn:aws:neptune-db:${AWS::Region}:${AWS::AccountId}:${ClusterResourceId}/*
|
||||
- ClusterResourceId: !GetAtt NeptuneCluster.ClusterResourceId
|
||||
|
||||
NeptuneUserAccessKey:
|
||||
Type: AWS::IAM::AccessKey
|
||||
Properties:
|
||||
UserName: !Ref NeptuneUser
|
||||
|
||||
# =============================================================================
|
||||
# OUTPUTS
|
||||
# =============================================================================
|
||||
|
||||
Outputs:
|
||||
NeptuneEndpoint:
|
||||
Description: Neptune cluster endpoint (hostname only) - use as NEPTUNE_ENDPOINT
|
||||
Value: !GetAtt NeptuneCluster.Endpoint
|
||||
|
||||
NeptunePort:
|
||||
Description: Neptune cluster port - use as NEPTUNE_PORT
|
||||
Value: !GetAtt NeptuneCluster.Port
|
||||
|
||||
AwsAccessKeyId:
|
||||
Description: Access key ID for the Neptune IAM user - use as AWS_ACCESS_KEY_ID
|
||||
Value: !Ref NeptuneUserAccessKey
|
||||
|
||||
AwsSecretAccessKey:
|
||||
Description: Secret access key for the Neptune IAM user - use as AWS_SECRET_ACCESS_KEY
|
||||
Value: !GetAtt NeptuneUserAccessKey.SecretAccessKey
|
||||
|
||||
AwsRegion:
|
||||
Description: AWS region where Neptune is deployed - use as AWS_REGION
|
||||
Value: !Ref AWS::Region
|
||||
|
||||
NeptuneClusterResourceId:
|
||||
Description: Neptune cluster resource ID (for IAM policy reference)
|
||||
Value: !GetAtt NeptuneCluster.ClusterResourceId
|
||||
|
||||
VpcId:
|
||||
Description: VPC ID
|
||||
Value: !Ref VPC
|
||||
|
||||
SecurityGroupId:
|
||||
Description: Neptune security group ID
|
||||
Value: !Ref NeptuneSecurityGroup
|
||||
@@ -110,7 +110,7 @@
|
||||
"source": [
|
||||
"# Set up API keys\n",
|
||||
"# Note: In production, use environment variables: export GROQ_API_KEY=\"your-key\"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"Your Groq API\")\n"
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -30,7 +30,7 @@
|
||||
"# Environment Setup\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ['GROQ_API_KEY'] = os.getenv('GROQ_API_KEY', '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')\n",
|
||||
"os.environ['GROQ_API_KEY'] = os.getenv('GROQ_API_KEY', '')\n",
|
||||
"\n",
|
||||
"# Install Semantica and all required dependencies\n",
|
||||
"%pip install -qU semantica networkx matplotlib plotly pandas faiss-cpu beautifulsoup4 groq sentence-transformers\n"
|
||||
@@ -84,7 +84,7 @@
|
||||
"source": [
|
||||
"# Set up API keys\n",
|
||||
"# Note: In production, use environment variables: export GROQ_API_KEY=\"your-key\"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"your-groq-api-key-here\")\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n",
|
||||
"\n",
|
||||
"print(\"API keys configured.\")\n"
|
||||
]
|
||||
|
||||
@@ -109,7 +109,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n"
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -85,7 +85,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n"
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -81,7 +81,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n",
|
||||
"\n",
|
||||
"# Configuration constants\n",
|
||||
"EMBEDDING_DIMENSION = 384\n",
|
||||
|
||||
@@ -98,7 +98,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n",
|
||||
"\n",
|
||||
"# Configuration constants\n",
|
||||
"EMBEDDING_DIMENSION = 384\n",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -83,7 +83,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n",
|
||||
"\n",
|
||||
"# Configuration constants\n",
|
||||
"EMBEDDING_DIMENSION = 384\n",
|
||||
|
||||
@@ -80,7 +80,7 @@
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"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\")\n",
|
||||
"os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n",
|
||||
"\n",
|
||||
"# Configuration constants\n",
|
||||
"EMBEDDING_DIMENSION = 384\n",
|
||||
|
||||
+5
-5
@@ -12,22 +12,22 @@ How to cite Semantica in academic papers and research.
|
||||
author = {Hawksight AI},
|
||||
year = {2026},
|
||||
url = {https://github.com/Hawksight-AI/semantica},
|
||||
version = {0.1.1},
|
||||
version = {0.2.3},
|
||||
doi = {10.5281/zenodo.XXXXXXX}
|
||||
}
|
||||
```
|
||||
|
||||
### APA
|
||||
Hawksight AI. (2026). *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering* (Version 0.1.1) [Computer software]. https://github.com/Hawksight-AI/semantica
|
||||
Hawksight AI. (2026). *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering* (Version 0.2.3) [Computer software]. https://github.com/Hawksight-AI/semantica
|
||||
|
||||
### MLA
|
||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.1.1, GitHub, 2026, https://github.com/Hawksight-AI/semantica.
|
||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.2.3, GitHub, 2026, https://github.com/Hawksight-AI/semantica.
|
||||
|
||||
### Chicago
|
||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.1.1. GitHub, 2026. https://github.com/Hawksight-AI/semantica.
|
||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.2.3. GitHub, 2026. https://github.com/Hawksight-AI/semantica.
|
||||
|
||||
### IEEE
|
||||
Hawksight AI, "Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering," Version 0.1.1, GitHub, 2026. [Online]. Available: https://github.com/Hawksight-AI/semantica
|
||||
Hawksight AI, "Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering," Version 0.2.3, GitHub, 2026. [Online]. Available: https://github.com/Hawksight-AI/semantica
|
||||
|
||||
---
|
||||
|
||||
|
||||
+1
-1
@@ -17,7 +17,7 @@
|
||||
|
||||
<p><em>The missing fabric between raw data and AI engineering. A comprehensive open-source framework for building semantic layers and knowledge engineering systems that transform unstructured data into AI-ready knowledge — powering Knowledge Graph-Powered RAG (GraphRAG), AI Agents, Multi-Agent Systems, and AI applications with structured semantic knowledge.</em></p>
|
||||
|
||||
<p>🆓 <strong>100% Open Source</strong> • 📜 <strong>MIT Licensed</strong> • 🚀 <strong>Latest Version: 0.1.1</strong> • 🚀 <strong>Production Ready</strong> • 🌍 <strong>Community Driven</strong></p>
|
||||
<p>🆓 <strong>100% Open Source</strong> • 📜 <strong>MIT Licensed</strong> • 🚀 <strong>Latest Version: 0.2.3</strong> • 🚀 <strong>Production Ready</strong> • 🌍 <strong>Community Driven</strong></p>
|
||||
|
||||
<p>
|
||||
<a href="getting-started/" class="md-button md-button--primary">Get Started</a>
|
||||
|
||||
@@ -9,7 +9,11 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
|
||||
// Define versions
|
||||
var versions = [
|
||||
{ name: "0.1.1", url: "#", current: true },
|
||||
{ name: "0.2.3", url: "#", current: true },
|
||||
{ name: "0.2.2", url: "#", current: false },
|
||||
{ name: "0.2.1", url: "#", current: false },
|
||||
{ name: "0.2.0", url: "#", current: false },
|
||||
{ name: "0.1.1", url: "#", current: false },
|
||||
{ name: "0.1.0", url: "#", current: false }
|
||||
];
|
||||
|
||||
|
||||
@@ -23,7 +23,7 @@ The **Semantic Extract Module** extracts structured information from unstructure
|
||||
- **High Accuracy**: LLM-based extraction for complex schemas
|
||||
- **Flexible Configuration**: Customize extraction for your domain
|
||||
- **Confidence Scores**: Get confidence scores for all extractions
|
||||
- **Batch Processing**: Efficient batch processing for large datasets
|
||||
- **Batch Processing**: Efficient parallel batch processing for large datasets
|
||||
- **Coreference Resolution**: Resolve pronouns to their entity references
|
||||
|
||||
### How It Works
|
||||
@@ -185,7 +185,9 @@ Core entity extraction implementation used by notebooks and lower-level integrat
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str or list | `"ml"` | Method(s): "ml", "llm", "pattern", "regex", "huggingface" |
|
||||
| `silent_fail` | bool | `False` | Return empty list on error instead of raising (LLM only) |
|
||||
| `max_text_length` | int | `None` | Max text length for auto-chunking (LLM only) |
|
||||
| `max_text_length` | int | `64000` | Max text length for auto-chunking (LLM only) |
|
||||
| `max_tokens` | int | `None` | Max output tokens for LLM generation |
|
||||
| `max_workers` | int | `1` | Threads for parallel batch processing |
|
||||
| `**config` | dict | `{}` | Method-specific config (e.g., `model`, `provider`) |
|
||||
|
||||
**Methods:**
|
||||
@@ -204,11 +206,12 @@ from semantica.semantic_extract import NERExtractor
|
||||
extractor = NERExtractor(method="ml", model="en_core_web_trf")
|
||||
entities = extractor.extract("Elon Musk leads SpaceX.")
|
||||
|
||||
# 2. LLM (OpenAI/Gemini/etc)
|
||||
# 2. LLM (OpenAI/Gemini/Groq/etc)
|
||||
extractor = NERExtractor(
|
||||
method="llm",
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
provider="groq",
|
||||
model="llama-3.3-70b-versatile",
|
||||
max_tokens=2048, # Increased output limit
|
||||
temperature=0.0
|
||||
)
|
||||
|
||||
@@ -232,6 +235,7 @@ Extracts relationships between entities.
|
||||
| `bidirectional` | bool | `False` | Extract bidirectional relations |
|
||||
| `confidence_threshold` | float | `0.6` | Minimum confidence score |
|
||||
| `max_distance` | int | `50` | Max token distance between entities |
|
||||
| `max_workers` | int | `1` | Threads for parallel batch processing |
|
||||
|
||||
**Methods:**
|
||||
|
||||
@@ -253,7 +257,7 @@ entities = ner.extract_entities(text)
|
||||
# Basic relation extraction
|
||||
rel_extractor = RelationExtractor()
|
||||
relations = rel_extractor.extract(text, entities=entities)
|
||||
# [Relation(source="Elon Musk", target="SpaceX", type="founded")]
|
||||
# [Relation(subject="Elon Musk", predicate="founded", object="SpaceX")]
|
||||
|
||||
# With configuration
|
||||
rel_extractor = RelationExtractor(
|
||||
@@ -308,6 +312,7 @@ Identifies events with temporal information and participants.
|
||||
| `extract_participants` | bool | `True` | Extract event participants |
|
||||
| `extract_location` | bool | `True` | Extract event locations |
|
||||
| `extract_time` | bool | `True` | Extract temporal information |
|
||||
| `max_workers` | int | `1` | Threads for parallel batch processing |
|
||||
|
||||
**Methods:**
|
||||
|
||||
@@ -340,7 +345,9 @@ Extracts RDF triplets (Subject-Predicate-Object).
|
||||
| `include_provenance` | bool | `False` | Track source sentences |
|
||||
| `method` | str | `"pattern"` | Extraction method ("pattern", "rules", "huggingface", "llm") |
|
||||
| `silent_fail` | bool | `False` | Return empty list on error instead of raising (LLM only) |
|
||||
| `max_text_length` | int | `None` | Max text length for auto-chunking (LLM only) |
|
||||
| `max_text_length` | int | `64000` | Max text length for auto-chunking (LLM only) |
|
||||
| `max_tokens` | int | `None` | Max output tokens for LLM generation |
|
||||
| `max_workers` | int | `1` | Threads for parallel batch processing |
|
||||
|
||||
**Methods:**
|
||||
|
||||
@@ -371,6 +378,7 @@ Extracts structured semantic networks with nodes and edges.
|
||||
|-----------|------|---------|-------------|
|
||||
| `ner_method` | str | `None` | Method for node extraction |
|
||||
| `relation_method` | str | `None` | Method for edge extraction |
|
||||
| `max_workers` | int | `1` | Threads for parallel batch processing |
|
||||
| `**config` | dict | `{}` | Configuration for underlying extractors |
|
||||
|
||||
**Methods:**
|
||||
@@ -424,6 +432,44 @@ enhanced_entities = extractor.enhance_entities(text, entities)
|
||||
|
||||
---
|
||||
|
||||
## Batch Processing & Provenance
|
||||
|
||||
All extractors support batch processing for high-throughput extraction. You can pass a list of strings or a list of dictionaries (with `content` and `id` keys).
|
||||
|
||||
**Features:**
|
||||
- **Progress Tracking**: Automatically shows a progress bar for large batches.
|
||||
- **Provenance Metadata**: Each extracted item includes `batch_index` and `document_id` in its `metadata`.
|
||||
|
||||
```python
|
||||
from semantica.semantic_extract import NERExtractor
|
||||
|
||||
documents = [
|
||||
{"id": "doc_1", "content": "Apple Inc. was founded by Steve Jobs."},
|
||||
{"id": "doc_2", "content": "Microsoft Corporation was founded by Bill Gates."}
|
||||
]
|
||||
|
||||
extractor = NERExtractor()
|
||||
batch_results = extractor.extract(documents)
|
||||
|
||||
for i, doc_entities in enumerate(batch_results):
|
||||
print(f"Document {i} entities:")
|
||||
for entity in doc_entities:
|
||||
print(f" - {entity.text} ({entity.label})")
|
||||
print(f" Provenance: Batch Index {entity.metadata['batch_index']}, Doc ID {entity.metadata.get('document_id')}")
|
||||
```
|
||||
|
||||
## Robust Extraction Fallbacks
|
||||
|
||||
The framework implements robust fallback chains to prevent empty results when primary methods fail (e.g., due to model unavailability or obscure text).
|
||||
|
||||
- **NER**: `ML/LLM` -> `Pattern` -> `Last Resort` (Capitalized Words)
|
||||
- **Relation**: `Primary` -> `Pattern` -> `Last Resort` (Adjacency)
|
||||
- **Triplet**: `Primary` -> `Relation-to-Triplet` -> `Pattern`
|
||||
|
||||
This ensures that you almost always get *some* structured data, even if it requires falling back to simpler heuristics.
|
||||
|
||||
---
|
||||
|
||||
## Usage Examples
|
||||
|
||||
```python
|
||||
|
||||
@@ -112,6 +112,8 @@ The main facade for all vector operations.
|
||||
| Method | Description |
|
||||
|--------|-------------|
|
||||
| `store_vectors(vectors, metadata)` | Store embeddings |
|
||||
| `add_documents(documents, metadata, batch_size, parallel)` | **(New)** Store documents with automatic embedding generation and parallelization |
|
||||
| `embed_batch(texts)` | **(New)** Generate embeddings for a batch of texts |
|
||||
| `search(query, k)` | Semantic search |
|
||||
| `delete(ids)` | Remove vectors |
|
||||
|
||||
@@ -120,15 +122,24 @@ The main facade for all vector operations.
|
||||
```python
|
||||
from semantica.vector_store import VectorStore
|
||||
|
||||
# Initialize (defaults to FAISS)
|
||||
# Initialize (defaults to FAISS, parallel enabled by default with 6 workers)
|
||||
store = VectorStore(backend="faiss", dimension=1536)
|
||||
|
||||
# Store
|
||||
# 1. Store pre-computed vectors
|
||||
ids = store.store_vectors(
|
||||
vectors=[[0.1, 0.2, ...], ...],
|
||||
metadata=[{"text": "Hello"}, ...]
|
||||
)
|
||||
|
||||
# 2. Store raw documents (High Performance)
|
||||
# Automatically handles embedding generation in parallel batches (uses default 6 workers)
|
||||
ids = store.add_documents(
|
||||
documents=["Doc 1", "Doc 2", ...],
|
||||
metadata=[{"id": 1}, {"id": 2}, ...],
|
||||
batch_size=32,
|
||||
parallel=True
|
||||
)
|
||||
|
||||
# Search
|
||||
results = store.search(query_vector=[0.1, 0.2, ...], k=5)
|
||||
```
|
||||
@@ -771,6 +782,7 @@ print(f"Context: {context}")
|
||||
---
|
||||
|
||||
## See Also
|
||||
- [High-Performance Usage Guide](../vector_store_usage.md) - **(New)** Parallel ingestion and batching guide
|
||||
- [Embeddings Module](embeddings.md) - Generates the vectors
|
||||
- [Context Module](context.md) - Uses vector store for memory
|
||||
- [Ingest Module](ingest.md) - Source of data
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
# High-Performance Vector Store Usage
|
||||
|
||||
This guide demonstrates how to leverage the new high-performance features of the Semantica Vector Store, specifically designed for efficient batch processing and parallel ingestion of large document sets.
|
||||
|
||||
## 🚀 Key Features
|
||||
|
||||
- **Parallel Ingestion**: Utilize multi-threading to embed and store documents concurrently.
|
||||
- **Batch Processing**: Automatically group documents into batches to minimize overhead.
|
||||
- **Unified API**: A single `add_documents` method handles embedding generation and storage.
|
||||
|
||||
---
|
||||
|
||||
## ⚡ Quick Start: Parallel Ingestion
|
||||
|
||||
The fastest way to ingest documents is using the `add_documents` method. Parallelization is enabled by default with optimized settings (6 workers).
|
||||
|
||||
```python
|
||||
from semantica.vector_store import VectorStore
|
||||
import time
|
||||
|
||||
store = VectorStore(
|
||||
backend="faiss",
|
||||
dimension=768,
|
||||
)
|
||||
|
||||
documents = [f"This is document number {i} with some content." for i in range(1000)]
|
||||
metadata = [{"source": "generated", "id": i} for i in range(1000)]
|
||||
|
||||
start_time = time.time()
|
||||
ids = store.add_documents(
|
||||
documents=documents,
|
||||
metadata=metadata,
|
||||
batch_size=64,
|
||||
parallel=True,
|
||||
)
|
||||
print(f"Ingested {len(ids)} documents in {time.time() - start_time:.2f}s")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Comparison
|
||||
|
||||
### Old Method (Sequential Loop)
|
||||
*Slower due to sequential processing and overhead per single item.*
|
||||
|
||||
```python
|
||||
for doc in documents:
|
||||
emb = embedder.generate(doc)
|
||||
store.store_vectors([emb], [{"text": doc}])
|
||||
```
|
||||
|
||||
### New Method (Parallel Batching)
|
||||
*Significantly faster (3x-10x) by utilizing thread pools and batch operations.*
|
||||
|
||||
```python
|
||||
store.add_documents(documents, parallel=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🛠 Configuration & Tuning
|
||||
|
||||
### `max_workers`
|
||||
Controls the number of concurrent threads used for embedding generation.
|
||||
- **Default**: 6 (Optimized for most systems)
|
||||
- **Recommendation**: You generally don't need to change this. If you have very high core counts or specific throughput needs, you can override it.
|
||||
|
||||
```python
|
||||
store = VectorStore(max_workers=16)
|
||||
```
|
||||
|
||||
### `batch_size`
|
||||
Controls how many documents are processed in a single chunk.
|
||||
- **Default**: 32
|
||||
- **Recommendation**:
|
||||
- **Local Models**: 32-64 usually works well.
|
||||
- **API Models (OpenAI, etc.)**: Larger batches (e.g., 100-200) can reduce network latency overhead.
|
||||
|
||||
```python
|
||||
store.add_documents(documents, batch_size=100)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🧩 Advanced: Manual Batch Embedding
|
||||
|
||||
If you need the embeddings without storing them immediately, use `embed_batch`.
|
||||
|
||||
```python
|
||||
vectors = store.embed_batch(
|
||||
texts=documents[:100],
|
||||
)
|
||||
|
||||
print(f"Generated {len(vectors)} vectors")
|
||||
```
|
||||
|
||||
## ⚠️ Best Practices
|
||||
|
||||
1. **Metadata Consistency**: Ensure your `metadata` list has the same length as your `documents` list.
|
||||
2. **Error Handling**: The `add_documents` method will propagate exceptions if embedding fails. Ensure your data is clean.
|
||||
3. **Memory Usage**: Very large `batch_size` combined with high `max_workers` can increase memory usage. Monitor your system resources.
|
||||
+165
-276
@@ -4,309 +4,198 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "semantica"
|
||||
version = "0.1.1"
|
||||
description = "🧠 Semantica - An Open Source Framework for building Semantic Layers and Knowledge Engineering "
|
||||
version = "0.2.3"
|
||||
description = "🧠 Semantica - An Open Source Framework for building Semantic Layers and Knowledge Engineering"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
authors = [
|
||||
{name = "Hawksight AI", email = "semantica-dev@users.noreply.github.com"}
|
||||
]
|
||||
maintainers = [
|
||||
{name = "Hawksight AI", email = "semantica-dev@users.noreply.github.com"}
|
||||
]
|
||||
license = { text = "MIT" }
|
||||
|
||||
authors = [{ name = "Hawksight AI", email = "semantica-dev@users.noreply.github.com" }]
|
||||
maintainers = [{ name = "Hawksight AI", email = "semantica-dev@users.noreply.github.com" }]
|
||||
|
||||
requires-python = ">=3.8"
|
||||
|
||||
classifiers = [
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Science/Research",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
"Topic :: Text Processing :: Linguistic",
|
||||
"Topic :: Database :: Database Engines/Servers",
|
||||
"Topic :: Internet :: WWW/HTTP :: Indexing/Search"
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Science/Research",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules"
|
||||
]
|
||||
|
||||
keywords = [
|
||||
"semantic-layer", "knowledge-engineering", "nlp", "knowledge-graph",
|
||||
"embeddings", "entity-extraction", "relationship-extraction", "rdf",
|
||||
"ontology", "semantic-analysis", "ai", "machine-learning"
|
||||
"semantic-layer", "knowledge-graph", "nlp", "embeddings",
|
||||
"entity-extraction", "relationship-extraction", "rdf", "ontology"
|
||||
]
|
||||
|
||||
# ---------------- CORE DEPENDENCIES (SAFE DEFAULT) ----------------
|
||||
dependencies = [
|
||||
"numpy>=1.21.0",
|
||||
"pandas>=1.3.0",
|
||||
"scikit-learn>=1.0.0",
|
||||
"umap-learn>=0.5.0",
|
||||
"spacy>=3.4.0",
|
||||
"transformers>=4.20.0",
|
||||
"torch>=1.12.0",
|
||||
"sentence-transformers>=2.2.0",
|
||||
"rdflib>=6.2.0",
|
||||
"networkx>=2.8.0",
|
||||
"matplotlib>=3.5.0",
|
||||
"seaborn>=0.11.0",
|
||||
"plotly>=5.10.0",
|
||||
"ipywidgets>=8.0.0",
|
||||
"requests>=2.28.0",
|
||||
"GitPython>=3.1.30",
|
||||
"chardet>=5.1.0",
|
||||
"protobuf==4.25.8",
|
||||
"grpcio==1.67.1",
|
||||
"beautifulsoup4>=4.11.0",
|
||||
"lxml>=4.9.0",
|
||||
"pypdf2>=2.10.0",
|
||||
"python-docx>=0.8.11",
|
||||
"docling>=1.0.0",
|
||||
"openpyxl>=3.0.10",
|
||||
"pillow>=9.2.0",
|
||||
"librosa>=0.9.0",
|
||||
"opencv-python>=4.6.0",
|
||||
"faiss-cpu>=1.7.0",
|
||||
"fastembed>=0.2.0",
|
||||
"onnxruntime>=1.17.0",
|
||||
"tokenizers>=0.15.0",
|
||||
"weaviate-client>=3.15.0",
|
||||
"qdrant-client>=1.3.0",
|
||||
"neo4j>=5.0.0",
|
||||
"falkordb>=1.0.0",
|
||||
"pymongo>=4.2.0",
|
||||
"sqlalchemy>=1.4.0",
|
||||
"psycopg2-binary>=2.9.0",
|
||||
"pymysql>=1.0.0",
|
||||
"redis>=4.3.0",
|
||||
"celery>=5.2.0",
|
||||
"kafka-python>=2.0.0",
|
||||
"pulsar-client>=3.0.0",
|
||||
"pika>=1.3.0",
|
||||
"boto3>=1.24.0",
|
||||
"azure-storage-blob>=12.12.0",
|
||||
"google-cloud-storage>=2.5.0",
|
||||
"pydantic>=2.0.0",
|
||||
"fastmcp>=0.1.0",
|
||||
"groq>=0.4.0",
|
||||
"openai>=1.0.0",
|
||||
"litellm>=1.0.0",
|
||||
"click>=8.1.0",
|
||||
"rich>=12.5.0",
|
||||
"tqdm>=4.64.0",
|
||||
"pyyaml>=6.0",
|
||||
"toml>=0.10.0",
|
||||
"python-dotenv>=0.20.0",
|
||||
"loguru>=0.6.0",
|
||||
"structlog>=22.1.0",
|
||||
"prometheus-client>=0.14.0",
|
||||
"opentelemetry-api>=1.12.0",
|
||||
"opentelemetry-sdk>=1.12.0",
|
||||
"opentelemetry-instrumentation",
|
||||
"fastapi>=0.78.0",
|
||||
"uvicorn>=0.18.0",
|
||||
"pytest>=7.1.0",
|
||||
"pytest-cov>=3.0.0",
|
||||
"pytest-asyncio>=0.19.0",
|
||||
"black>=22.6.0",
|
||||
"isort>=5.10.0",
|
||||
"flake8>=4.0.0",
|
||||
"mypy>=0.971",
|
||||
"pre-commit>=2.19.0"
|
||||
"numpy>=1.21.0",
|
||||
"pandas>=1.3.0",
|
||||
"scikit-learn>=1.0.0",
|
||||
"umap-learn>=0.5.0",
|
||||
"spacy>=3.4.0",
|
||||
"transformers>=4.20.0",
|
||||
"torch>=1.12.0",
|
||||
"sentence-transformers>=2.2.0",
|
||||
"rdflib>=6.2.0",
|
||||
"networkx>=2.8.0",
|
||||
"matplotlib>=3.5.0",
|
||||
"seaborn>=0.11.0",
|
||||
"plotly>=5.10.0",
|
||||
"ipywidgets>=8.0.0",
|
||||
"requests>=2.28.0",
|
||||
"GitPython>=3.1.30",
|
||||
"chardet>=5.1.0",
|
||||
"protobuf>=5.29.1,<7.0",
|
||||
"grpcio>=1.71.2",
|
||||
"beautifulsoup4>=4.11.0",
|
||||
"lxml>=4.9.0",
|
||||
"pypdf2>=2.10.0",
|
||||
"python-docx>=0.8.11",
|
||||
"openpyxl>=3.0.10",
|
||||
"pillow>=9.2.0",
|
||||
"librosa>=0.9.0",
|
||||
"opencv-python>=4.6.0",
|
||||
"faiss-cpu>=1.7.0",
|
||||
"fastembed>=0.2.0",
|
||||
"onnxruntime>=1.17.0",
|
||||
"tokenizers>=0.15.0",
|
||||
"pydantic>=2.0.0",
|
||||
"click>=8.1.0",
|
||||
"rich>=12.5.0",
|
||||
"tqdm>=4.64.0",
|
||||
"pyyaml>=6.0",
|
||||
"toml>=0.10.0",
|
||||
"python-dotenv>=0.20.0",
|
||||
"loguru>=0.6.0",
|
||||
"structlog>=22.1.0"
|
||||
]
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/Hawksight-AI/semantica"
|
||||
Repository = "https://github.com/Hawksight-AI/semantica"
|
||||
"Bug Tracker" = "https://github.com/Hawksight-AI/semantica/issues"
|
||||
Discussions = "https://github.com/Hawksight-AI/semantica/discussions"
|
||||
|
||||
# ---------------- OPTIONAL DEPENDENCIES ----------------
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=7.1.0",
|
||||
"pytest-cov>=3.0.0",
|
||||
"pytest-asyncio>=0.19.0",
|
||||
"black>=22.6.0",
|
||||
"isort>=5.10.0",
|
||||
"flake8>=4.0.0",
|
||||
"mypy>=0.971",
|
||||
"pre-commit>=2.19.0",
|
||||
"jupyter>=1.0.0",
|
||||
"ipykernel>=6.15.0",
|
||||
"notebook>=6.4.0"
|
||||
]
|
||||
viz = [
|
||||
"pyvis>=0.3.0",
|
||||
"graphviz>=0.20.0",
|
||||
"umap-learn>=0.5.0",
|
||||
"d3blocks>=1.0.0"
|
||||
]
|
||||
gpu = [
|
||||
"torch>=1.12.0",
|
||||
"faiss-gpu>=1.7.0",
|
||||
"cupy>=10.0.0"
|
||||
]
|
||||
cloud = [
|
||||
"boto3>=1.24.0",
|
||||
"azure-storage-blob>=12.12.0",
|
||||
"google-cloud-storage>=2.5.0",
|
||||
"kubernetes>=24.0.0",
|
||||
"helm>=3.10.0"
|
||||
]
|
||||
monitoring = [
|
||||
"prometheus-client>=0.14.0",
|
||||
"opentelemetry-api>=1.12.0",
|
||||
"opentelemetry-sdk>=1.12.0",
|
||||
"opentelemetry-instrumentation>=0.32.0",
|
||||
"grafana-api>=1.0.0",
|
||||
"elasticsearch>=8.5.0"
|
||||
]
|
||||
llm-openai = [
|
||||
"openai>=1.0.0"
|
||||
]
|
||||
llm-gemini = [
|
||||
"google-generativeai>=0.3.0"
|
||||
]
|
||||
llm-groq = [
|
||||
"groq>=0.4.0"
|
||||
]
|
||||
llm-anthropic = [
|
||||
"anthropic>=0.18.0"
|
||||
]
|
||||
llm-ollama = [
|
||||
"ollama>=0.1.0"
|
||||
]
|
||||
llm-deepseek = [
|
||||
"deepseek>=0.1.0"
|
||||
]
|
||||
llm-litellm = [
|
||||
"litellm>=1.0.0"
|
||||
]
|
||||
|
||||
# ---- LLM Providers ----
|
||||
llm-openai = ["openai>=1.0.0"]
|
||||
llm-groq = ["groq>=0.4.0"]
|
||||
llm-gemini = ["google-genai>=0.1.0"]
|
||||
llm-anthropic = ["anthropic>=0.18.0"]
|
||||
llm-ollama = ["ollama>=0.1.0"]
|
||||
llm-deepseek = ["deepseek>=0.1.0"]
|
||||
llm-litellm = ["litellm>=1.0.0"]
|
||||
llm-instructor = ["instructor>=1.0.0"]
|
||||
|
||||
llm-all = [
|
||||
"semantica[llm-openai,llm-gemini,llm-groq,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm]"
|
||||
]
|
||||
models-huggingface = [
|
||||
"transformers>=4.20.0",
|
||||
"torch>=1.12.0"
|
||||
]
|
||||
split-tiktoken = [
|
||||
"tiktoken>=0.5.0"
|
||||
]
|
||||
split-community = [
|
||||
"python-louvain>=0.16"
|
||||
]
|
||||
split-topic = [
|
||||
"bertopic>=0.15.0",
|
||||
"gensim>=4.3.0"
|
||||
]
|
||||
split-all = [
|
||||
"semantica[split-tiktoken,split-community,split-topic]"
|
||||
]
|
||||
graph-neo4j = [
|
||||
"neo4j>=5.0.0"
|
||||
]
|
||||
graph-falkordb = [
|
||||
"falkordb>=1.0.0",
|
||||
"redis>=4.3.0"
|
||||
]
|
||||
graph-all = [
|
||||
"semantica[graph-neo4j,graph-falkordb]"
|
||||
]
|
||||
parse-docling = [
|
||||
"docling>=1.0.0"
|
||||
]
|
||||
all = [
|
||||
"semantica[dev,viz,gpu,cloud,monitoring,llm-all,models-huggingface,split-all,graph-all,parse-docling]"
|
||||
"semantica[llm-openai,llm-groq,llm-gemini,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm,llm-instructor]"
|
||||
]
|
||||
|
||||
# ---- Document Parsing ----
|
||||
parse-docling = ["docling>=1.0.0"]
|
||||
|
||||
# ---- Embedding / Models ----
|
||||
models-huggingface = [
|
||||
"transformers>=4.20.0",
|
||||
"torch>=1.12.0"
|
||||
]
|
||||
|
||||
# ---- Graph Backends ----
|
||||
graph-neo4j = ["neo4j>=5.0.0"]
|
||||
graph-falkordb = ["falkordb>=1.0.0", "redis>=4.3.0"]
|
||||
graph-amazon-neptune = ["boto3>=1.24.0", "neo4j>=5.0.0"]
|
||||
|
||||
graph-all = [
|
||||
"semantica[graph-neo4j,graph-falkordb,graph-amazon-neptune]"
|
||||
]
|
||||
|
||||
# ---- Infra / Queues / Workers ----
|
||||
infra = [
|
||||
"redis>=4.3.0",
|
||||
"celery>=5.2.0",
|
||||
"kafka-python>=2.0.0",
|
||||
"pulsar-client>=3.0.0",
|
||||
"pika>=1.3.0"
|
||||
]
|
||||
|
||||
# ---- Cloud Providers ----
|
||||
cloud = [
|
||||
"boto3>=1.24.0",
|
||||
"azure-storage-blob>=12.12.0",
|
||||
"google-cloud-storage>=2.5.0"
|
||||
]
|
||||
|
||||
# ---- Monitoring (FIXED) ----
|
||||
monitoring = [
|
||||
"prometheus-client>=0.14.0",
|
||||
"opentelemetry-api>=1.30.0,<2.0.0",
|
||||
"opentelemetry-sdk>=1.30.0,<2.0.0",
|
||||
"opentelemetry-semantic-conventions>=0.58b0,<0.61b0",
|
||||
"opentelemetry-instrumentation>=0.58b0,<0.61b0"
|
||||
]
|
||||
|
||||
# ---- Visualization ----
|
||||
viz = [
|
||||
"pyvis>=0.3.0",
|
||||
"graphviz>=0.20.0",
|
||||
"d3blocks>=1.0.0"
|
||||
]
|
||||
|
||||
# ---- GPU ----
|
||||
gpu = [
|
||||
"faiss-gpu>=1.7.0",
|
||||
"cupy>=10.0.0"
|
||||
]
|
||||
|
||||
# ---- Splitting / Chunking ----
|
||||
split-tiktoken = ["tiktoken>=0.5.0"]
|
||||
split-community = ["python-louvain>=0.16"]
|
||||
split-topic = ["bertopic>=0.15.0", "gensim>=4.3.0"]
|
||||
|
||||
split-all = [
|
||||
"semantica[split-tiktoken,split-community,split-topic]"
|
||||
]
|
||||
|
||||
# ---- Dev ----
|
||||
dev = [
|
||||
"pytest>=7.1.0",
|
||||
"pytest-cov>=3.0.0",
|
||||
"pytest-asyncio>=0.19.0",
|
||||
"black>=22.6.0",
|
||||
"isort>=5.10.0",
|
||||
"flake8>=4.0.0",
|
||||
"mypy>=0.971",
|
||||
"pre-commit>=2.19.0",
|
||||
"jupyter>=1.0.0",
|
||||
"ipykernel>=6.15.0"
|
||||
]
|
||||
|
||||
# ---- Everything ----
|
||||
all = [
|
||||
"semantica[dev,viz,gpu,infra,cloud,monitoring,llm-all,models-huggingface,split-all,graph-all,parse-docling]"
|
||||
]
|
||||
|
||||
# ---------------- ENTRYPOINTS ----------------
|
||||
[project.scripts]
|
||||
semantica = "semantica.cli:main"
|
||||
semantica-server = "semantica.server:main"
|
||||
semantica-worker = "semantica.worker:main"
|
||||
|
||||
# ---------------- TOOLING ----------------
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["semantica*"]
|
||||
exclude = ["tests*", "docs*", "examples*"]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
semantica = ["*.yaml", "*.yml", "*.json", "*.toml", "*.txt", "*.md"]
|
||||
|
||||
[tool.black]
|
||||
line-length = 88
|
||||
target-version = ['py38', 'py39', 'py310', 'py311', 'py312']
|
||||
include = '\.pyi?$'
|
||||
extend-exclude = '''
|
||||
/(
|
||||
# directories
|
||||
\.eggs
|
||||
| \.git
|
||||
| \.hg
|
||||
| \.mypy_cache
|
||||
| \.tox
|
||||
| \.venv
|
||||
| build
|
||||
| dist
|
||||
)/
|
||||
'''
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
multi_line_output = 3
|
||||
line_length = 88
|
||||
known_first_party = ["semantica"]
|
||||
known_third_party = ["numpy", "pandas", "scikit-learn", "spacy", "transformers", "torch"]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.9"
|
||||
warn_return_any = true
|
||||
warn_unused_configs = true
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
check_untyped_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
no_implicit_optional = true
|
||||
warn_redundant_casts = true
|
||||
warn_unused_ignores = true
|
||||
warn_no_return = true
|
||||
warn_unreachable = true
|
||||
strict_equality = true
|
||||
show_error_codes = true
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
minversion = "7.0"
|
||||
addopts = "-ra -q --strict-markers --strict-config"
|
||||
testpaths = ["tests"]
|
||||
python_files = ["test_*.py", "*_test.py"]
|
||||
python_classes = ["Test*"]
|
||||
python_functions = ["test_*"]
|
||||
markers = [
|
||||
"slow: marks tests as slow (deselect with '-m \"not slow\"')",
|
||||
"integration: marks tests as integration tests",
|
||||
"unit: marks tests as unit tests",
|
||||
"gpu: marks tests that require GPU",
|
||||
"cloud: marks tests that require cloud services"
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
source = ["semantica"]
|
||||
omit = [
|
||||
"*/tests/*",
|
||||
"*/test_*",
|
||||
"*/__pycache__/*",
|
||||
"*/migrations/*"
|
||||
]
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_lines = [
|
||||
"pragma: no cover",
|
||||
"def __repr__",
|
||||
"if self.debug:",
|
||||
"if settings.DEBUG",
|
||||
"raise AssertionError",
|
||||
"raise NotImplementedError",
|
||||
"if 0:",
|
||||
"if __name__ == .__main__.:",
|
||||
"class .*\\bProtocol\\):",
|
||||
"@(abc\\.)?abstractmethod"
|
||||
]
|
||||
|
||||
@@ -10,7 +10,7 @@ Main exports:
|
||||
- Config: Configuration management
|
||||
"""
|
||||
|
||||
__version__ = "0.1.1"
|
||||
__version__ = "0.2.3"
|
||||
__author__ = "Semantica Contributors"
|
||||
__license__ = "MIT"
|
||||
|
||||
|
||||
@@ -1,48 +1,71 @@
|
||||
"""
|
||||
Graph Store Module
|
||||
|
||||
This module provides comprehensive property graph database integration for the
|
||||
Semantica framework, supporting multiple graph database backends including Neo4j
|
||||
and FalkorDB for storing and querying knowledge graphs.
|
||||
This module provides comprehensive property graph database integration for
|
||||
the Semantica framework, supporting multiple graph database backends including
|
||||
Neo4j and FalkorDB for storing and querying knowledge graphs.
|
||||
|
||||
Algorithms Used:
|
||||
|
||||
Graph Store Management:
|
||||
- Store Registration: Store type detection, store factory pattern, configuration management, default store selection
|
||||
- Backend Pattern: Unified interface for multiple backends (Neo4j, FalkorDB), backend instantiation, backend-specific operation delegation
|
||||
- Store Selection: Default store resolution, store ID lookup, store validation
|
||||
- Store Registration: Store type detection, store factory pattern,
|
||||
configuration management, default store selection
|
||||
- Backend Pattern: Unified interface for multiple backends (Neo4j,
|
||||
FalkorDB), backend instantiation, backend-specific operation delegation
|
||||
- Store Selection: Default store resolution, store ID lookup,
|
||||
store validation
|
||||
|
||||
Node and Relationship Operations:
|
||||
- Node Creation: Single node insertion, batch node insertion, property validation, label management, backend delegation
|
||||
- Node Retrieval: Pattern matching (label/property filtering), Cypher query construction, result extraction, node reconstruction
|
||||
- Node Update: Property update, label modification, atomic update operations, conflict detection
|
||||
- Node Deletion: Node matching, cascade deletion (optional), deletion operation delegation, result verification
|
||||
- Relationship Creation: Single relationship insertion, batch insertion, property validation, type management
|
||||
- Node Creation: Single node insertion, batch node insertion,
|
||||
property validation, label management, backend delegation
|
||||
- Node Retrieval: Pattern matching (label/property filtering),
|
||||
Cypher query construction, result extraction, node reconstruction
|
||||
- Node Update: Property update, label modification, atomic update
|
||||
operations, conflict detection
|
||||
- Node Deletion: Node matching, cascade deletion (optional),
|
||||
deletion operation delegation, result verification
|
||||
- Relationship Creation: Single relationship insertion, batch insertion,
|
||||
property validation, type management
|
||||
- Relationship Retrieval: Pattern matching, path queries, traversal queries
|
||||
- Relationship Update: Property update, type modification
|
||||
- Relationship Deletion: Relationship matching, deletion operation delegation
|
||||
- Relationship Deletion: Relationship matching, deletion operation
|
||||
delegation
|
||||
|
||||
Graph Query Execution:
|
||||
- Cypher Query: Full Cypher query language support for Neo4j and FalkorDB (OpenCypher)
|
||||
- Pattern Matching: Node and relationship pattern matching, variable binding, path matching
|
||||
- Graph Traversal: BFS/DFS traversal, shortest path algorithms, path finding
|
||||
- Cypher Query: Full Cypher query language support for Neo4j and
|
||||
FalkorDB (OpenCypher)
|
||||
- Pattern Matching: Node and relationship pattern matching, variable
|
||||
binding, path matching
|
||||
- Graph Traversal: BFS/DFS traversal, shortest path algorithms,
|
||||
path finding
|
||||
- Aggregation: COUNT, SUM, AVG, MIN, MAX operations, GROUP BY support
|
||||
- Query Optimization: Query caching, execution plan analysis, index utilization
|
||||
- Query Optimization: Query caching, execution plan analysis,
|
||||
index utilization
|
||||
|
||||
Graph Analytics:
|
||||
- Centrality Algorithms: Degree centrality, betweenness centrality, PageRank, closeness centrality
|
||||
- Community Detection: Label propagation, Louvain modularity, connected components
|
||||
- Path Algorithms: Shortest path, all shortest paths, Dijkstra, A* pathfinding
|
||||
- Centrality Algorithms: Degree centrality, betweenness centrality,
|
||||
PageRank, closeness centrality
|
||||
- Community Detection: Label propagation, Louvain modularity,
|
||||
connected components
|
||||
- Path Algorithms: Shortest path, all shortest paths, Dijkstra,
|
||||
A* pathfinding
|
||||
- Similarity: Node similarity, Jaccard similarity, cosine similarity
|
||||
|
||||
Store Backends:
|
||||
- Neo4j Store: Official Neo4j Python driver, Bolt protocol communication, transaction support, multi-database support, APOC procedures
|
||||
- FalkorDB Store: Redis-based graph database, sparse matrix representation, linear algebra queries, OpenCypher support, ultra-fast performance
|
||||
- Neo4j Store: Official Neo4j Python driver, Bolt protocol
|
||||
communication, transaction support, multi-database support,
|
||||
APOC procedures
|
||||
- FalkorDB Store: Redis-based graph database, sparse matrix
|
||||
representation, linear algebra queries, OpenCypher support,
|
||||
ultra-fast performance
|
||||
|
||||
Bulk Operations:
|
||||
- Batch Processing: Chunking algorithm (fixed-size batch creation), batch size optimization, memory management for large datasets
|
||||
- Transaction Management: ACID transaction support, batch commits, rollback on failure
|
||||
- Progress Tracking: Load progress calculation, elapsed time tracking, throughput calculation
|
||||
- Batch Processing: Chunking algorithm (fixed-size batch creation),
|
||||
batch size optimization, memory management for large datasets
|
||||
- Transaction Management: ACID transaction support, batch commits,
|
||||
rollback on failure
|
||||
- Progress Tracking: Load progress calculation, elapsed time tracking,
|
||||
throughput calculation
|
||||
|
||||
Key Features:
|
||||
- Multi-backend property graph support (Neo4j, FalkorDB)
|
||||
@@ -79,33 +102,42 @@ Convenience Functions:
|
||||
- list_available_methods: List registered graph store methods
|
||||
|
||||
Example Usage:
|
||||
>>> from semantica.graph_store import GraphStore, create_node, create_relationship, execute_query
|
||||
>>> from semantica.graph_store import GraphStore, create_node, \
|
||||
... create_relationship, execute_query
|
||||
>>> # Using convenience functions
|
||||
>>> node_id = create_node(labels=["Person"], properties={"name": "Alice", "age": 30})
|
||||
>>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id, rel_type="KNOWS", properties={"since": 2020})
|
||||
>>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 RETURN p.name")
|
||||
>>> node_id = create_node(labels=["Person"],
|
||||
... properties={"name": "Alice", "age": 30})
|
||||
>>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id,
|
||||
... rel_type="KNOWS",
|
||||
... properties={"since": 2020})
|
||||
>>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 "
|
||||
... "RETURN p.name")
|
||||
>>> # Using classes directly
|
||||
>>> store = GraphStore(backend="neo4j", uri="bolt://localhost:7687")
|
||||
>>> node_id = store.create_node(labels=["Person"], properties={"name": "Bob"})
|
||||
>>> node_id = store.create_node(labels=["Person"],
|
||||
... properties={"name": "Bob"})
|
||||
>>> results = store.execute_query("MATCH (n) RETURN n LIMIT 10")
|
||||
|
||||
Author: Semantica Contributors
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
from .config import GraphStoreConfig, graph_store_config
|
||||
from .falkordb_store import (
|
||||
FalkorDBStore,
|
||||
FalkorDBClient,
|
||||
FalkorDBGraph,
|
||||
from .amazon_neptune import (
|
||||
AmazonNeptuneStore,
|
||||
NeptuneAuthTokenManager,
|
||||
NeptuneDriver,
|
||||
NeptuneSession,
|
||||
NeptuneTransaction,
|
||||
)
|
||||
from .config import GraphStoreConfig, graph_store_config
|
||||
from .falkordb_store import FalkorDBClient, FalkorDBGraph, FalkorDBStore
|
||||
from .graph_store import (
|
||||
GraphAnalytics,
|
||||
GraphManager,
|
||||
GraphStore,
|
||||
NodeManager,
|
||||
QueryEngine,
|
||||
RelationshipManager,
|
||||
GraphAnalytics,
|
||||
)
|
||||
from .methods import (
|
||||
create_node,
|
||||
@@ -125,11 +157,7 @@ from .methods import (
|
||||
update_node,
|
||||
update_relationship,
|
||||
)
|
||||
from .neo4j_store import (
|
||||
Neo4jStore,
|
||||
Neo4jDriver,
|
||||
Neo4jTransaction,
|
||||
)
|
||||
from .neo4j_store import Neo4jDriver, Neo4jStore, Neo4jTransaction
|
||||
from .registry import MethodRegistry, method_registry
|
||||
|
||||
__all__ = [
|
||||
@@ -144,6 +172,12 @@ __all__ = [
|
||||
"Neo4jStore",
|
||||
"Neo4jDriver",
|
||||
"Neo4jTransaction",
|
||||
# Amazon Neptune
|
||||
"AmazonNeptuneStore",
|
||||
"NeptuneAuthTokenManager",
|
||||
"NeptuneDriver",
|
||||
"NeptuneSession",
|
||||
"NeptuneTransaction",
|
||||
# FalkorDB
|
||||
"FalkorDBStore",
|
||||
"FalkorDBClient",
|
||||
@@ -171,4 +205,3 @@ __all__ = [
|
||||
"MethodRegistry",
|
||||
"method_registry",
|
||||
]
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,7 +6,8 @@ supporting multiple configuration sources including environment variables, confi
|
||||
and programmatic configuration.
|
||||
|
||||
Supported Configuration Sources:
|
||||
- Environment variables: GRAPH_STORE_DEFAULT_BACKEND, GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc.
|
||||
- Environment variables: GRAPH_STORE_DEFAULT_BACKEND,
|
||||
GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc.
|
||||
- Config files: YAML, JSON, TOML formats
|
||||
- Programmatic: Python API for setting graph store configurations
|
||||
|
||||
@@ -44,7 +45,11 @@ from ..utils.logging import get_logger
|
||||
|
||||
|
||||
class GraphStoreConfig:
|
||||
"""Configuration manager for graph store module - supports .env files, environment variables, and programmatic config."""
|
||||
"""
|
||||
Configuration manager for graph store module.
|
||||
|
||||
Supports .env files, environment variables, and programmatic config.
|
||||
"""
|
||||
|
||||
def __init__(self, config_file: Optional[str] = None):
|
||||
"""
|
||||
@@ -124,6 +129,15 @@ class GraphStoreConfig:
|
||||
"GRAPH_STORE_FALKORDB_PORT": "falkordb_port",
|
||||
"GRAPH_STORE_FALKORDB_PASSWORD": "falkordb_password",
|
||||
"GRAPH_STORE_FALKORDB_GRAPH_NAME": "falkordb_graph_name",
|
||||
# Amazon Neptune settings
|
||||
"GRAPH_STORE_NEPTUNE_ENDPOINT": "neptune_endpoint",
|
||||
"GRAPH_STORE_NEPTUNE_PORT": "neptune_port",
|
||||
"GRAPH_STORE_NEPTUNE_REGION": "neptune_region",
|
||||
"GRAPH_STORE_NEPTUNE_IAM_AUTH": "neptune_iam_auth",
|
||||
"GRAPH_STORE_NEPTUNE_USE_SSL": "neptune_use_ssl",
|
||||
"AWS_ACCESS_KEY_ID": "neptune_access_key",
|
||||
"AWS_SECRET_ACCESS_KEY": "neptune_secret_key",
|
||||
"AWS_SESSION_TOKEN": "neptune_session_token",
|
||||
}
|
||||
|
||||
for env_var, config_key in env_mappings.items():
|
||||
@@ -135,6 +149,7 @@ class GraphStoreConfig:
|
||||
"timeout",
|
||||
"max_retries",
|
||||
"falkordb_port",
|
||||
"neptune_port",
|
||||
]:
|
||||
try:
|
||||
self._config[config_key] = int(value)
|
||||
@@ -142,7 +157,11 @@ class GraphStoreConfig:
|
||||
self.logger.warning(
|
||||
f"Invalid integer value for {env_var}: {value}"
|
||||
)
|
||||
elif config_key in ["neo4j_encrypted"]:
|
||||
elif config_key in [
|
||||
"neo4j_encrypted",
|
||||
"neptune_iam_auth",
|
||||
"neptune_use_ssl",
|
||||
]:
|
||||
self._config[config_key] = value.lower() in [
|
||||
"true",
|
||||
"1",
|
||||
@@ -171,6 +190,15 @@ class GraphStoreConfig:
|
||||
"falkordb_port": 6379,
|
||||
"falkordb_password": None,
|
||||
"falkordb_graph_name": "default",
|
||||
# Amazon Neptune defaults
|
||||
"neptune_endpoint": None,
|
||||
"neptune_port": 8182,
|
||||
"neptune_region": None,
|
||||
"neptune_iam_auth": True,
|
||||
"neptune_use_ssl": True,
|
||||
"neptune_access_key": None,
|
||||
"neptune_secret_key": None,
|
||||
"neptune_session_token": None,
|
||||
}
|
||||
|
||||
for key, default_value in defaults.items():
|
||||
@@ -269,6 +297,24 @@ class GraphStoreConfig:
|
||||
"graph_name": self._config.get("falkordb_graph_name"),
|
||||
}
|
||||
|
||||
def get_neptune_config(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get Amazon Neptune-specific configuration.
|
||||
|
||||
Returns:
|
||||
Neptune configuration dictionary
|
||||
"""
|
||||
return {
|
||||
"endpoint": self._config.get("neptune_endpoint"),
|
||||
"port": self._config.get("neptune_port"),
|
||||
"region": self._config.get("neptune_region"),
|
||||
"iam_auth": self._config.get("neptune_iam_auth"),
|
||||
"use_ssl": self._config.get("neptune_use_ssl"),
|
||||
"access_key": self._config.get("neptune_access_key"),
|
||||
"secret_key": self._config.get("neptune_secret_key"),
|
||||
"session_token": self._config.get("neptune_session_token"),
|
||||
}
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset configuration to defaults."""
|
||||
self._config.clear()
|
||||
@@ -278,4 +324,3 @@ class GraphStoreConfig:
|
||||
|
||||
# Global configuration instance
|
||||
graph_store_config = GraphStoreConfig()
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ License: MIT
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.exceptions import ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
from ..utils.progress_tracker import get_progress_tracker
|
||||
from .config import graph_store_config
|
||||
@@ -214,7 +214,9 @@ class RelationshipManager:
|
||||
Returns:
|
||||
List of relationships
|
||||
"""
|
||||
return self.backend.get_relationships(node_id, rel_type, direction, limit, **options)
|
||||
return self.backend.get_relationships(
|
||||
node_id, rel_type, direction, limit, **options
|
||||
)
|
||||
|
||||
def delete(
|
||||
self,
|
||||
@@ -290,6 +292,7 @@ class QueryEngine:
|
||||
) -> str:
|
||||
"""Generate cache key for query."""
|
||||
import hashlib
|
||||
|
||||
key_str = f"{query}:{str(parameters)}"
|
||||
return hashlib.md5(key_str.encode()).hexdigest()
|
||||
|
||||
@@ -340,7 +343,9 @@ class GraphAnalytics:
|
||||
Returns:
|
||||
Path information or None
|
||||
"""
|
||||
return self.backend.shortest_path(start_node_id, end_node_id, rel_type, max_depth, **options)
|
||||
return self.backend.shortest_path(
|
||||
start_node_id, end_node_id, rel_type, max_depth, **options
|
||||
)
|
||||
|
||||
def get_neighbors(
|
||||
self,
|
||||
@@ -363,7 +368,9 @@ class GraphAnalytics:
|
||||
Returns:
|
||||
List of neighboring nodes
|
||||
"""
|
||||
return self.backend.get_neighbors(node_id, rel_type, direction, depth, **options)
|
||||
return self.backend.get_neighbors(
|
||||
node_id, rel_type, direction, depth, **options
|
||||
)
|
||||
|
||||
def degree_centrality(
|
||||
self,
|
||||
@@ -440,7 +447,7 @@ class GraphAnalytics:
|
||||
Component information
|
||||
"""
|
||||
backend_type = type(self.backend).__name__
|
||||
|
||||
|
||||
if "Neo4j" in backend_type:
|
||||
query = """
|
||||
CALL gds.wcc.stream({
|
||||
@@ -452,16 +459,23 @@ class GraphAnalytics:
|
||||
"""
|
||||
params = {"label": labels[0] if labels else "*"}
|
||||
result = self.backend.execute_query(query, params)
|
||||
return [{"component": r["componentId"], "nodes": r["nodes"]} for r in result]
|
||||
|
||||
return [
|
||||
{"component": r["componentId"], "nodes": r["nodes"]} for r in result
|
||||
]
|
||||
|
||||
elif "NetworkX" in backend_type:
|
||||
import networkx as nx
|
||||
|
||||
G = self.backend.graph
|
||||
components = list(nx.connected_components(G))
|
||||
return [{"component": i, "nodes": list(c)} for i, c in enumerate(components)]
|
||||
|
||||
return [
|
||||
{"component": i, "nodes": list(c)} for i, c in enumerate(components)
|
||||
]
|
||||
|
||||
else:
|
||||
raise NotImplementedError(f"connected_components not implemented for {backend_type}")
|
||||
raise NotImplementedError(
|
||||
f"connected_components not implemented for {backend_type}"
|
||||
)
|
||||
|
||||
|
||||
class GraphManager:
|
||||
@@ -534,7 +548,11 @@ class GraphStore:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
# Determine backend
|
||||
self.backend = backend or config.get("backend") or graph_store_config.get("default_backend", "neo4j")
|
||||
self.backend = (
|
||||
backend
|
||||
or config.get("backend")
|
||||
or graph_store_config.get("default_backend", "neo4j")
|
||||
)
|
||||
self.config = config
|
||||
|
||||
# Initialize store backend
|
||||
@@ -546,16 +564,25 @@ class GraphStore:
|
||||
"""Initialize the appropriate store backend based on backend type."""
|
||||
if self.backend == "neo4j":
|
||||
from .neo4j_store import Neo4jStore
|
||||
|
||||
neo4j_config = graph_store_config.get_neo4j_config()
|
||||
neo4j_config.update(self.config)
|
||||
self._store_backend = Neo4jStore(**neo4j_config)
|
||||
|
||||
elif self.backend == "falkordb":
|
||||
from .falkordb_store import FalkorDBStore
|
||||
|
||||
falkordb_config = graph_store_config.get_falkordb_config()
|
||||
falkordb_config.update(self.config)
|
||||
self._store_backend = FalkorDBStore(**falkordb_config)
|
||||
|
||||
elif self.backend == "neptune" or self.backend == "amazon_neptune":
|
||||
from .amazon_neptune import AmazonNeptuneStore
|
||||
|
||||
neptune_config = graph_store_config.get_neptune_config()
|
||||
neptune_config.update(self.config)
|
||||
self._store_backend = AmazonNeptuneStore(**neptune_config)
|
||||
|
||||
else:
|
||||
raise ValidationError(f"Unknown backend: {self.backend}")
|
||||
|
||||
@@ -621,7 +648,9 @@ class GraphStore:
|
||||
**options,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get nodes matching criteria."""
|
||||
return self._manager.nodes.get(labels=labels, properties=properties, limit=limit, **options)
|
||||
return self._manager.nodes.get(
|
||||
labels=labels, properties=properties, limit=limit, **options
|
||||
)
|
||||
|
||||
def update_node(
|
||||
self,
|
||||
@@ -665,7 +694,9 @@ class GraphStore:
|
||||
**options,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get relationships."""
|
||||
return self._manager.relationships.get(node_id, rel_type, direction, limit, **options)
|
||||
return self._manager.relationships.get(
|
||||
node_id, rel_type, direction, limit, **options
|
||||
)
|
||||
|
||||
def delete_relationship(
|
||||
self,
|
||||
@@ -723,7 +754,9 @@ class GraphStore:
|
||||
node_id, rel_type, direction, actual_depth, **options
|
||||
)
|
||||
|
||||
def query(self, query: str, parameters: Optional[Dict[str, Any]] = None, **options) -> List[Dict[str, Any]]:
|
||||
def query(
|
||||
self, query: str, parameters: Optional[Dict[str, Any]] = None, **options
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute a query and return results (Compatibility method for ContextRetriever).
|
||||
|
||||
@@ -771,42 +804,46 @@ class GraphStore:
|
||||
# Convert to GraphStore format (labels, properties)
|
||||
graph_nodes = []
|
||||
for node in nodes:
|
||||
# Extract label from type
|
||||
labels = [node.get("type", "Entity")]
|
||||
if isinstance(labels[0], str):
|
||||
labels = [labels[0]] # Ensure list
|
||||
# Extract labels - support both 'labels' array and 'type' string
|
||||
labels = node.get("labels")
|
||||
if not labels:
|
||||
node_type = node.get("type", "Entity")
|
||||
labels = [node_type] if isinstance(node_type, str) else node_type
|
||||
if isinstance(labels, str):
|
||||
labels = [labels]
|
||||
|
||||
# Prepare properties
|
||||
props = node.get("properties", {}).copy()
|
||||
|
||||
|
||||
# Ensure ID is preserved
|
||||
if "id" in node and "id" not in props:
|
||||
props["id"] = node["id"]
|
||||
|
||||
|
||||
# Ensure content/text is preserved
|
||||
if "content" in node and "content" not in props:
|
||||
props["content"] = node["content"]
|
||||
if "text" in node and "text" not in props:
|
||||
props["text"] = node["text"]
|
||||
|
||||
graph_nodes.append({
|
||||
"labels": labels,
|
||||
"properties": props
|
||||
})
|
||||
graph_nodes.append({"labels": labels, "properties": props})
|
||||
|
||||
# Use batch creation
|
||||
# Note: create_nodes expects dicts with 'labels' and 'properties' keys if passed directly?
|
||||
# Note: create_nodes expects dicts with 'labels' and 'properties'
|
||||
# keys if passed directly?
|
||||
# Let's check create_nodes signature implementation in manager.
|
||||
# But here I'll assume create_nodes takes a list of such dicts or similar.
|
||||
# But here I'll assume create_nodes takes a list of such dicts
|
||||
# or similar.
|
||||
# Actually, let's look at create_nodes wrapper in this file:
|
||||
# def create_nodes(self, nodes: List[Dict[str, Any]], **options)
|
||||
# It passes to self._manager.nodes.create_batch(nodes)
|
||||
|
||||
|
||||
# If create_batch expects specific format, I should match it.
|
||||
# Assuming create_batch is smart enough or expects standard format.
|
||||
# To be safe, let's look at NodeManager.create_batch if possible, but I can't easily.
|
||||
# Standard expectation: List of dicts where each dict has labels and properties.
|
||||
|
||||
# To be safe, let's look at NodeManager.create_batch if possible,
|
||||
# but I can't easily.
|
||||
# Standard expectation: List of dicts where each dict has labels
|
||||
# and properties.
|
||||
|
||||
result = self.create_nodes(graph_nodes, **options)
|
||||
return len(result)
|
||||
|
||||
@@ -827,17 +864,21 @@ class GraphStore:
|
||||
target_id = edge.get("target_id")
|
||||
rel_type = edge.get("type", "RELATED_TO")
|
||||
properties = edge.get("properties", {}).copy()
|
||||
|
||||
|
||||
# Preserve weight
|
||||
if "weight" in edge:
|
||||
properties["weight"] = edge["weight"]
|
||||
|
||||
if source_id and target_id:
|
||||
try:
|
||||
self.create_relationship(source_id, target_id, rel_type, properties, **options)
|
||||
self.create_relationship(
|
||||
source_id, target_id, rel_type, properties, **options
|
||||
)
|
||||
count += 1
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to add edge {source_id}->{target_id}: {e}")
|
||||
self.logger.warning(
|
||||
f"Failed to add edge {source_id}->{target_id}: {e}"
|
||||
)
|
||||
return count
|
||||
|
||||
def build_from_conversations(
|
||||
@@ -872,27 +913,29 @@ class GraphStore:
|
||||
all_nodes = []
|
||||
all_edges = []
|
||||
seen_nodes = set()
|
||||
|
||||
|
||||
for conv in conversations:
|
||||
# Load conversation if string (file path)
|
||||
conv_data = conv
|
||||
if isinstance(conv, str):
|
||||
from pathlib import Path
|
||||
|
||||
from ..utils.helpers import read_json_file
|
||||
|
||||
conv_data = read_json_file(Path(conv))
|
||||
|
||||
nodes, edges = self._process_conversation_to_elements(
|
||||
conv_data,
|
||||
conv_data,
|
||||
extract_intents=extract_intents,
|
||||
extract_sentiments=extract_sentiments
|
||||
extract_sentiments=extract_sentiments,
|
||||
)
|
||||
|
||||
|
||||
# Add unique nodes
|
||||
for node in nodes:
|
||||
if node["id"] not in seen_nodes:
|
||||
all_nodes.append(node)
|
||||
seen_nodes.add(node["id"])
|
||||
|
||||
|
||||
all_edges.extend(edges)
|
||||
|
||||
if link_entities:
|
||||
@@ -904,13 +947,8 @@ class GraphStore:
|
||||
edge_count = self.add_edges(all_edges)
|
||||
|
||||
self.progress_tracker.stop_tracking(tracking_id, status="completed")
|
||||
|
||||
return {
|
||||
"statistics": {
|
||||
"node_count": node_count,
|
||||
"edge_count": edge_count
|
||||
}
|
||||
}
|
||||
|
||||
return {"statistics": {"node_count": node_count, "edge_count": edge_count}}
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
@@ -930,92 +968,112 @@ class GraphStore:
|
||||
"""
|
||||
nodes = []
|
||||
edges = []
|
||||
|
||||
|
||||
# Process entities
|
||||
for entity in entities:
|
||||
entity_id = entity.get("id") or entity.get("entity_id")
|
||||
if entity_id:
|
||||
nodes.append({
|
||||
"id": entity_id,
|
||||
"type": entity.get("type", "entity"),
|
||||
"properties": {
|
||||
"content": entity.get("text") or entity.get("label") or entity_id,
|
||||
**entity
|
||||
nodes.append(
|
||||
{
|
||||
"id": entity_id,
|
||||
"type": entity.get("type", "entity"),
|
||||
"properties": {
|
||||
"content": entity.get("text")
|
||||
or entity.get("label")
|
||||
or entity_id,
|
||||
**entity,
|
||||
},
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
# Process relationships
|
||||
for rel in relationships:
|
||||
source = rel.get("source_id")
|
||||
target = rel.get("target_id")
|
||||
if source and target:
|
||||
edges.append({
|
||||
"source_id": source,
|
||||
"target_id": target,
|
||||
"type": rel.get("type", "related_to"),
|
||||
"weight": rel.get("confidence", 1.0),
|
||||
"properties": rel
|
||||
})
|
||||
edges.append(
|
||||
{
|
||||
"source_id": source,
|
||||
"target_id": target,
|
||||
"type": rel.get("type", "related_to"),
|
||||
"weight": rel.get("confidence", 1.0),
|
||||
"properties": rel,
|
||||
}
|
||||
)
|
||||
|
||||
node_count = self.add_nodes(nodes)
|
||||
edge_count = self.add_edges(edges)
|
||||
|
||||
|
||||
return {"statistics": {"node_count": node_count, "edge_count": edge_count}}
|
||||
|
||||
def _process_conversation_to_elements(self, conv_data: Dict[str, Any], **kwargs) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||
def _process_conversation_to_elements(
|
||||
self, conv_data: Dict[str, Any], **kwargs
|
||||
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||
"""Helper to process conversation into nodes and edges."""
|
||||
nodes = []
|
||||
edges = []
|
||||
|
||||
|
||||
conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}"
|
||||
|
||||
# Conversation node
|
||||
nodes.append({
|
||||
"id": conv_id,
|
||||
"type": "conversation",
|
||||
"properties": {
|
||||
"content": conv_data.get("content", "") or conv_data.get("summary", ""),
|
||||
"timestamp": conv_data.get("timestamp")
|
||||
nodes.append(
|
||||
{
|
||||
"id": conv_id,
|
||||
"type": "conversation",
|
||||
"properties": {
|
||||
"content": conv_data.get("content", "")
|
||||
or conv_data.get("summary", ""),
|
||||
"timestamp": conv_data.get("timestamp"),
|
||||
},
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
name_to_id = {}
|
||||
extract_entities = kwargs.get("extract_entities", True) # Default true if not passed?
|
||||
# Actually ContextGraph defaults to True in init, but here we are static.
|
||||
# Note: extract_entities option is available but not used in this
|
||||
# implementation. Default true if not passed. ContextGraph defaults
|
||||
# to True in init, but here we are static.
|
||||
# Let's assume True unless told otherwise or check config.
|
||||
|
||||
|
||||
# Extract entities
|
||||
for entity in conv_data.get("entities", []):
|
||||
entity_id = entity.get("id") or entity.get("entity_id")
|
||||
entity_text = entity.get("text") or entity.get("label") or entity.get("name") or entity_id
|
||||
entity_text = (
|
||||
entity.get("text")
|
||||
or entity.get("label")
|
||||
or entity.get("name")
|
||||
or entity_id
|
||||
)
|
||||
entity_type = entity.get("type", "entity")
|
||||
|
||||
# Generate ID if missing
|
||||
if not entity_id and entity_text:
|
||||
import hashlib
|
||||
entity_hash = hashlib.md5(f"{entity_text}_{entity_type}".encode()).hexdigest()[:12]
|
||||
|
||||
entity_hash = hashlib.md5(
|
||||
f"{entity_text}_{entity_type}".encode()
|
||||
).hexdigest()[:12]
|
||||
entity_id = f"{entity_type.lower()}_{entity_hash}"
|
||||
|
||||
if entity_id:
|
||||
if entity_text:
|
||||
name_to_id[entity_text] = entity_id
|
||||
|
||||
nodes.append({
|
||||
"id": entity_id,
|
||||
"type": "entity", # Normalize type?
|
||||
"properties": {
|
||||
"content": entity_text,
|
||||
"type": entity_type,
|
||||
**entity
|
||||
nodes.append(
|
||||
{
|
||||
"id": entity_id,
|
||||
"type": "entity", # Normalize type?
|
||||
"properties": {
|
||||
"content": entity_text,
|
||||
"type": entity_type,
|
||||
**entity,
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
)
|
||||
|
||||
# Edge: Conversation -> Entity
|
||||
edges.append({
|
||||
"source_id": conv_id,
|
||||
"target_id": entity_id,
|
||||
"type": "mentions"
|
||||
})
|
||||
edges.append(
|
||||
{"source_id": conv_id, "target_id": entity_id, "type": "mentions"}
|
||||
)
|
||||
|
||||
# Extract relationships
|
||||
for rel in conv_data.get("relationships", []):
|
||||
@@ -1029,43 +1087,54 @@ class GraphStore:
|
||||
target = name_to_id[rel.get("target")]
|
||||
|
||||
if source and target:
|
||||
edges.append({
|
||||
"source_id": source,
|
||||
"target_id": target,
|
||||
"type": rel.get("type", "related_to"),
|
||||
"weight": rel.get("confidence", 1.0),
|
||||
"properties": rel
|
||||
})
|
||||
|
||||
edges.append(
|
||||
{
|
||||
"source_id": source,
|
||||
"target_id": target,
|
||||
"type": rel.get("type", "related_to"),
|
||||
"weight": rel.get("confidence", 1.0),
|
||||
"properties": rel,
|
||||
}
|
||||
)
|
||||
|
||||
return nodes, edges
|
||||
|
||||
def _link_entities_elements(self, nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
def _link_entities_elements(
|
||||
self, nodes: List[Dict[str, Any]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Link similar entities."""
|
||||
edges = []
|
||||
# Lazy import to avoid circular dependency
|
||||
try:
|
||||
from ..context.entity_linker import EntityLinker
|
||||
linker = EntityLinker() # Use default config
|
||||
|
||||
linker = EntityLinker() # Use default config
|
||||
except (ImportError, OSError):
|
||||
return []
|
||||
|
||||
entity_nodes = [n for n in nodes if n.get("type") == "entity"]
|
||||
for i, node1 in enumerate(entity_nodes):
|
||||
content1 = node1["properties"].get("content", "")
|
||||
if not content1: continue
|
||||
|
||||
if not content1:
|
||||
continue
|
||||
|
||||
for node2 in entity_nodes[i + 1 :]:
|
||||
content2 = node2["properties"].get("content", "")
|
||||
if not content2: continue
|
||||
|
||||
similarity = linker._calculate_text_similarity(content1.lower(), content2.lower())
|
||||
if not content2:
|
||||
continue
|
||||
|
||||
similarity = linker._calculate_text_similarity(
|
||||
content1.lower(), content2.lower()
|
||||
)
|
||||
if similarity >= linker.similarity_threshold:
|
||||
edges.append({
|
||||
"source_id": node1["id"],
|
||||
"target_id": node2["id"],
|
||||
"type": "similar_to",
|
||||
"weight": similarity
|
||||
})
|
||||
edges.append(
|
||||
{
|
||||
"source_id": node1["id"],
|
||||
"target_id": node2["id"],
|
||||
"type": "similar_to",
|
||||
"weight": similarity,
|
||||
}
|
||||
)
|
||||
return edges
|
||||
|
||||
@property
|
||||
@@ -1087,4 +1156,3 @@ class GraphStore:
|
||||
def analytics(self) -> GraphAnalytics:
|
||||
"""Get analytics engine."""
|
||||
return self._manager.analytics
|
||||
|
||||
|
||||
+134
-29
@@ -89,6 +89,7 @@ class GraphBuilder:
|
||||
self.track_history = track_history
|
||||
self.version_snapshots = version_snapshots
|
||||
self.graph_store = graph_store
|
||||
self.config = kwargs # Store additional config for extractors
|
||||
|
||||
# Initialize logging
|
||||
from ..utils.logging import get_logger
|
||||
@@ -130,6 +131,11 @@ class GraphBuilder:
|
||||
|
||||
def _process_item(self, item: Any, all_entities: List[Any], all_relationships: List[Any], **options):
|
||||
"""Helper to process a single item and add to entities or relationships list."""
|
||||
if isinstance(item, str):
|
||||
# Treat string as text for extraction
|
||||
self._extract_from_text(item, all_entities, all_relationships, **options)
|
||||
return
|
||||
|
||||
if hasattr(item, "text") and (hasattr(item, "label") or hasattr(item, "type")):
|
||||
# It's likely an Entity object
|
||||
entity_dict = {
|
||||
@@ -155,7 +161,14 @@ class GraphBuilder:
|
||||
}
|
||||
all_relationships.append(rel_dict)
|
||||
elif isinstance(item, dict):
|
||||
# Detect and normalize Entity objects inside dict
|
||||
if "source_id" in item and "source" not in item:
|
||||
item["source"] = item["source_id"]
|
||||
if "target_id" in item and "target" not in item:
|
||||
item["target"] = item["target_id"]
|
||||
if "subject" in item and "source" not in item:
|
||||
item["source"] = item["subject"]
|
||||
if "object" in item and "target" not in item:
|
||||
item["target"] = item["object"]
|
||||
if "source" in item and not isinstance(item["source"], str):
|
||||
src = item["source"]
|
||||
item["source"] = getattr(src, "id", getattr(src, "text", str(src)))
|
||||
@@ -209,30 +222,68 @@ class GraphBuilder:
|
||||
# If still nothing found and has 'text', try extraction
|
||||
if not found_something and "text" in item:
|
||||
text = item["text"]
|
||||
# Perform extraction if requested or if it's the only way
|
||||
if options.get("extract", True):
|
||||
from ..semantic_extract.ner_extractor import NERExtractor
|
||||
from ..semantic_extract.triplet_extractor import TripletExtractor
|
||||
|
||||
ner_method = options.get("ner_method", "ml")
|
||||
triplet_method = options.get("triplet_method", "pattern")
|
||||
|
||||
ner = NERExtractor(method=ner_method)
|
||||
entities = ner.extract_entities(text)
|
||||
for ent in entities:
|
||||
self._process_item(ent, all_entities, all_relationships, **options)
|
||||
|
||||
# Only try triplets if specifically requested or if method provided
|
||||
if "triplet_method" in options or options.get("extract_relations", False):
|
||||
triplet = TripletExtractor(method=triplet_method)
|
||||
relations = triplet.extract_triplets(text)
|
||||
for rel in relations:
|
||||
self._process_item(rel, all_entities, all_relationships, **options)
|
||||
found_something = True
|
||||
self._extract_from_text(text, all_entities, all_relationships, **options)
|
||||
found_something = True
|
||||
else:
|
||||
# Unknown type
|
||||
pass
|
||||
|
||||
def _extract_from_text(self, text: str, all_entities: List[Any], all_relationships: List[Any], **options):
|
||||
"""Helper to extract knowledge from text using configured methods."""
|
||||
if not options.get("extract", True):
|
||||
return
|
||||
|
||||
from ..semantic_extract.ner_extractor import NERExtractor
|
||||
from ..semantic_extract.relation_extractor import RelationExtractor
|
||||
from ..semantic_extract.triplet_extractor import TripletExtractor
|
||||
|
||||
# Default to LLM methods as per requirement
|
||||
ner_method = options.get("ner_method", "llm")
|
||||
relation_method = options.get("relation_method", "llm")
|
||||
triplet_method = options.get("triplet_method", "llm")
|
||||
|
||||
self.logger.info(f"Extracting knowledge from text ({len(text)} chars) using {ner_method}...")
|
||||
|
||||
# 1. Extract Entities
|
||||
ner = NERExtractor(method=ner_method, **self.config)
|
||||
try:
|
||||
entities = ner.extract_entities(text, **options)
|
||||
extracted_count = len(entities)
|
||||
self._extraction_stats["extracted_entities"] += extracted_count
|
||||
self.logger.info(f"Extracted {extracted_count} entities")
|
||||
for ent in entities:
|
||||
self._process_item(ent, all_entities, all_relationships, **options)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Entity extraction failed: {e}")
|
||||
entities = []
|
||||
|
||||
# 2. Extract Relations (if requested)
|
||||
if options.get("extract_relations", True):
|
||||
rel_extractor = RelationExtractor(method=relation_method, **self.config)
|
||||
try:
|
||||
# Pass entities if available to help relation extraction
|
||||
relations = rel_extractor.extract_relations(text, entities=entities, **options)
|
||||
extracted_count = len(relations)
|
||||
self._extraction_stats["extracted_relations"] += extracted_count
|
||||
self.logger.info(f"Extracted {extracted_count} relationships")
|
||||
for rel in relations:
|
||||
self._process_item(rel, all_entities, all_relationships, **options)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Relation extraction failed: {e}")
|
||||
|
||||
# 3. Extract Triplets (if requested)
|
||||
if options.get("extract_triplets", True):
|
||||
trip_extractor = TripletExtractor(method=triplet_method, **self.config)
|
||||
try:
|
||||
triplets = trip_extractor.extract_triplets(text, entities=entities, **options)
|
||||
extracted_count = len(triplets)
|
||||
self._extraction_stats["extracted_triplets"] += extracted_count
|
||||
self.logger.info(f"Extracted {extracted_count} triplets")
|
||||
for trip in triplets:
|
||||
self._process_item(trip, all_entities, all_relationships, **options)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Triplet extraction failed: {e}")
|
||||
|
||||
def build(
|
||||
self,
|
||||
sources: Union[List[Any], Any],
|
||||
@@ -303,8 +354,31 @@ class GraphBuilder:
|
||||
elif not isinstance(sources, list):
|
||||
sources = [sources]
|
||||
|
||||
# Count input relationships for warning if all are dropped
|
||||
input_relationships_count = 0
|
||||
if isinstance(source_dict, dict):
|
||||
rels = source_dict.get("relationships", [])
|
||||
if isinstance(rels, list):
|
||||
input_relationships_count += len(rels)
|
||||
elif rels is not None:
|
||||
input_relationships_count += 1
|
||||
if explicit_relationships:
|
||||
for rel_item in explicit_relationships:
|
||||
if isinstance(rel_item, list):
|
||||
input_relationships_count += len(rel_item)
|
||||
else:
|
||||
input_relationships_count += 1
|
||||
|
||||
# Track graph building
|
||||
build_start_time = time.time()
|
||||
|
||||
# Initialize extraction statistics for traceability
|
||||
self._extraction_stats = {
|
||||
"extracted_entities": 0,
|
||||
"extracted_relations": 0,
|
||||
"extracted_triplets": 0
|
||||
}
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="kg",
|
||||
submodule="GraphBuilder",
|
||||
@@ -416,11 +490,12 @@ class GraphBuilder:
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
# Check if relationships are already in dictionary format
|
||||
sample_rel = relationships_list[0] if relationships_list else None
|
||||
is_dict_format = isinstance(sample_rel, dict) and (
|
||||
"source" in sample_rel and "target" in sample_rel
|
||||
) and not hasattr(sample_rel, "__dict__") # Ensure it's not a class instance
|
||||
("source" in sample_rel and "target" in sample_rel)
|
||||
or ("source_id" in sample_rel and "target_id" in sample_rel)
|
||||
or ("subject" in sample_rel and "object" in sample_rel)
|
||||
) and not hasattr(sample_rel, "__dict__")
|
||||
|
||||
if is_dict_format:
|
||||
# Fast path: directly append dictionaries after normalizing source/target
|
||||
@@ -429,8 +504,15 @@ class GraphBuilder:
|
||||
batch = relationships_list[i:i + batch_size]
|
||||
for item in batch:
|
||||
if isinstance(item, dict):
|
||||
# Normalize source/target if they are objects
|
||||
rel_dict = item.copy()
|
||||
if "source_id" in rel_dict and "source" not in rel_dict:
|
||||
rel_dict["source"] = rel_dict["source_id"]
|
||||
if "target_id" in rel_dict and "target" not in rel_dict:
|
||||
rel_dict["target"] = rel_dict["target_id"]
|
||||
if "subject" in rel_dict and "source" not in rel_dict:
|
||||
rel_dict["source"] = rel_dict["subject"]
|
||||
if "object" in rel_dict and "target" not in rel_dict:
|
||||
rel_dict["target"] = rel_dict["object"]
|
||||
if "source" in rel_dict and not isinstance(rel_dict["source"], str):
|
||||
src = rel_dict["source"]
|
||||
rel_dict["source"] = getattr(src, "id", getattr(src, "text", str(src)))
|
||||
@@ -514,11 +596,19 @@ class GraphBuilder:
|
||||
resolution_start = time.time()
|
||||
resolved_entities = resolver_to_use.resolve_entities(all_entities)
|
||||
resolution_time = time.time() - resolution_start
|
||||
print(f"✅ Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)")
|
||||
print(f"[DONE] Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)")
|
||||
self.logger.info(
|
||||
f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities"
|
||||
)
|
||||
|
||||
if input_relationships_count > 0 and len(all_relationships) == 0:
|
||||
warning_msg = (
|
||||
f"All relationships were dropped during graph building: "
|
||||
f"{input_relationships_count} input relationships, 0 in final graph"
|
||||
)
|
||||
self.logger.warning(warning_msg)
|
||||
print(f"Warning: {warning_msg}")
|
||||
|
||||
# Build graph structure
|
||||
print("Building graph structure...")
|
||||
structure_start = time.time()
|
||||
@@ -534,7 +624,7 @@ class GraphBuilder:
|
||||
},
|
||||
}
|
||||
structure_time = time.time() - structure_start
|
||||
print(f"✅ Graph structure built ({structure_time:.2f}s)")
|
||||
print(f"[DONE] Graph structure built ({structure_time:.2f}s)")
|
||||
|
||||
# Persist to GraphStore if available
|
||||
if self.graph_store:
|
||||
@@ -567,7 +657,7 @@ class GraphBuilder:
|
||||
edge_time = time.time() - edge_start
|
||||
total_store_time = time.time() - store_start
|
||||
print(f" Added {edge_count} edges ({edge_time:.2f}s)")
|
||||
print(f"✅ GraphStore persistence complete ({total_store_time:.2f}s total)")
|
||||
print(f"[DONE] GraphStore persistence complete ({total_store_time:.2f}s total)")
|
||||
self.logger.info(f"Persisted {node_count} nodes and {edge_count} edges")
|
||||
|
||||
# Detect and resolve conflicts if conflict detector is available
|
||||
@@ -604,7 +694,14 @@ class GraphBuilder:
|
||||
|
||||
# Print final summary with timing
|
||||
print(f"\n{'='*60}")
|
||||
print(f"✅ Knowledge Graph Build Complete")
|
||||
print(f"[INFO] Extraction Statistics")
|
||||
print(f" Extracted Entities: {self._extraction_stats['extracted_entities']}")
|
||||
print(f" Extracted Relationships: {self._extraction_stats['extracted_relations']}")
|
||||
print(f" Extracted Triplets: {self._extraction_stats['extracted_triplets']}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"[DONE] Knowledge Graph Build Complete")
|
||||
print(f" Entities: {len(resolved_entities)}")
|
||||
print(f" Relationships: {len(all_relationships)}")
|
||||
print(f" Total time: {total_build_time:.2f}s")
|
||||
@@ -623,6 +720,14 @@ class GraphBuilder:
|
||||
)
|
||||
raise
|
||||
|
||||
def build_single_source(
|
||||
self,
|
||||
kg_data: Dict[str, Any],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Dict[str, Any]:
|
||||
return self.build(kg_data, pipeline_id=pipeline_id, **options)
|
||||
|
||||
def add_temporal_edge(
|
||||
self,
|
||||
graph,
|
||||
|
||||
@@ -32,12 +32,14 @@ License: MIT
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
from typing import Any, Callable, Dict, List, Optional, Union, TYPE_CHECKING
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
from ..utils.progress_tracker import get_progress_tracker
|
||||
from .pipeline_validator import PipelineValidator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .pipeline_validator import PipelineValidator
|
||||
|
||||
|
||||
class StepStatus(Enum):
|
||||
@@ -104,6 +106,8 @@ class PipelineBuilder:
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
from .pipeline_validator import PipelineValidator
|
||||
|
||||
self.validator = PipelineValidator(**self.config)
|
||||
self.steps: List[PipelineStep] = []
|
||||
self.step_registry: Dict[str, Callable] = {}
|
||||
|
||||
@@ -15,6 +15,8 @@ Key Features:
|
||||
- Semantic network construction
|
||||
- LLM-based extraction enhancement
|
||||
- Extraction validation and quality assessment
|
||||
- Batch processing with provenance tracking (batch_index, document_id)
|
||||
- Robust fallback mechanisms (ML -> Pattern -> Last Resort)
|
||||
|
||||
Main Classes:
|
||||
- NamedEntityRecognizer: Main NER coordinator (confidence_threshold, merge_overlapping)
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
Result Caching Module
|
||||
|
||||
This module provides caching mechanisms for extraction results to avoid redundant
|
||||
computations and API calls. It implements an LRU (Least Recently Used) cache
|
||||
with Time-To-Live (TTL) support.
|
||||
|
||||
Key Features:
|
||||
- LRU Caching: Evicts least recently used items when cache is full
|
||||
- TTL Support: Expires items after a configurable duration
|
||||
- Namespaced Caching: Separate caches for entities, relations, and triplets
|
||||
- Hash-based Keys: Uses stable hashing for text and parameters
|
||||
|
||||
Classes:
|
||||
- ExtractionCache: Main cache manager
|
||||
- CacheItem: Container for cached data with metadata
|
||||
|
||||
Author: Semantica Contributors
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
import time
|
||||
import hashlib
|
||||
import json
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Optional, Union, List
|
||||
from threading import Lock
|
||||
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
class CacheItem:
|
||||
"""Container for cached data."""
|
||||
def __init__(self, value: Any, ttl: Optional[int] = None):
|
||||
self.value = value
|
||||
self.timestamp = time.time()
|
||||
self.ttl = ttl
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if item has expired."""
|
||||
if self.ttl is None:
|
||||
return False
|
||||
return time.time() - self.timestamp > self.ttl
|
||||
|
||||
class ExtractionCache:
|
||||
"""
|
||||
LRU Cache for extraction results.
|
||||
Thread-safe implementation.
|
||||
"""
|
||||
def __init__(self, max_size: int = 1000, ttl: int = 3600):
|
||||
"""
|
||||
Initialize the cache.
|
||||
|
||||
Args:
|
||||
max_size: Maximum number of items to store per namespace
|
||||
ttl: Time to live in seconds (default 1 hour)
|
||||
"""
|
||||
self.max_size = max_size
|
||||
self.ttl = ttl
|
||||
self._caches: Dict[str, OrderedDict] = {
|
||||
"entities": OrderedDict(),
|
||||
"relations": OrderedDict(),
|
||||
"triplets": OrderedDict()
|
||||
}
|
||||
self._locks: Dict[str, Lock] = {
|
||||
"entities": Lock(),
|
||||
"relations": Lock(),
|
||||
"triplets": Lock()
|
||||
}
|
||||
self.logger = get_logger("extraction_cache")
|
||||
self.enabled = True
|
||||
|
||||
def _generate_key(self, text: str, **params) -> str:
|
||||
"""
|
||||
Generate a stable cache key based on text and parameters.
|
||||
|
||||
Note: Sensitive parameters like 'api_key' are excluded from the cache key
|
||||
to prevent security risks and ensure cache sharing where appropriate.
|
||||
"""
|
||||
# Filter out sensitive keys
|
||||
sensitive_keys = {'api_key', 'token', 'password', 'secret', 'auth', 'authorization'}
|
||||
filtered_params = {k: v for k, v in params.items() if k.lower() not in sensitive_keys}
|
||||
|
||||
# Create a stable string representation of params
|
||||
# Sort keys to ensure consistent ordering
|
||||
param_str = json.dumps(filtered_params, sort_keys=True, default=str)
|
||||
|
||||
# Combine text and params
|
||||
content = f"{text}|{param_str}"
|
||||
|
||||
# Return hash (SHA-256 for better security than MD5)
|
||||
return hashlib.sha256(content.encode('utf-8')).hexdigest()
|
||||
|
||||
def get(self, namespace: str, text: str, **params) -> Optional[Any]:
|
||||
"""
|
||||
Retrieve item from cache.
|
||||
|
||||
Args:
|
||||
namespace: Cache namespace ("entities", "relations", "triplets")
|
||||
text: Input text used for extraction
|
||||
**params: Extraction parameters used
|
||||
|
||||
Returns:
|
||||
Cached result or None if not found/expired
|
||||
"""
|
||||
if not self.enabled:
|
||||
return None
|
||||
|
||||
if namespace not in self._caches:
|
||||
return None
|
||||
|
||||
key = self._generate_key(text, **params)
|
||||
|
||||
with self._locks[namespace]:
|
||||
cache = self._caches[namespace]
|
||||
if key in cache:
|
||||
item = cache[key]
|
||||
|
||||
# Check expiration
|
||||
if item.is_expired():
|
||||
del cache[key]
|
||||
return None
|
||||
|
||||
# Move to end (mark as recently used)
|
||||
cache.move_to_end(key)
|
||||
return item.value
|
||||
|
||||
return None
|
||||
|
||||
def set(self, namespace: str, text: str, value: Any, **params) -> None:
|
||||
"""
|
||||
Add item to cache.
|
||||
|
||||
Args:
|
||||
namespace: Cache namespace
|
||||
text: Input text
|
||||
value: Result to cache
|
||||
**params: Extraction parameters
|
||||
"""
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
if namespace not in self._caches:
|
||||
self.logger.warning(f"Unknown cache namespace: {namespace}")
|
||||
return
|
||||
|
||||
key = self._generate_key(text, **params)
|
||||
item = CacheItem(value, self.ttl)
|
||||
|
||||
with self._locks[namespace]:
|
||||
cache = self._caches[namespace]
|
||||
|
||||
# If key exists, update and move to end
|
||||
if key in cache:
|
||||
cache.move_to_end(key)
|
||||
|
||||
cache[key] = item
|
||||
|
||||
# Evict if full
|
||||
if len(cache) > self.max_size:
|
||||
cache.popitem(last=False) # Remove first (least recently used)
|
||||
|
||||
def clear(self, namespace: Optional[str] = None):
|
||||
"""Clear cache(s)."""
|
||||
if namespace:
|
||||
if namespace in self._caches:
|
||||
with self._locks[namespace]:
|
||||
self._caches[namespace].clear()
|
||||
else:
|
||||
for ns in self._caches:
|
||||
with self._locks[ns]:
|
||||
self._caches[ns].clear()
|
||||
|
||||
def get_stats(self) -> Dict[str, Dict[str, int]]:
|
||||
"""Get cache statistics."""
|
||||
stats = {}
|
||||
for ns, cache in self._caches.items():
|
||||
stats[ns] = {
|
||||
"size": len(cache),
|
||||
"max_size": self.max_size
|
||||
}
|
||||
return stats
|
||||
|
||||
# Global cache instance
|
||||
extraction_cache = ExtractionCache()
|
||||
@@ -40,8 +40,9 @@ License: MIT
|
||||
"""
|
||||
|
||||
import os
|
||||
import multiprocessing
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
from typing import Dict, Optional, Any
|
||||
|
||||
from ..utils.logging import get_logger
|
||||
|
||||
@@ -53,9 +54,23 @@ class Config:
|
||||
"""Initialize configuration manager."""
|
||||
self.logger = get_logger("config")
|
||||
self._configs: Dict[str, Dict] = {}
|
||||
# Default optimization settings
|
||||
self._configs["optimization"] = {
|
||||
"enable_cache": True,
|
||||
"cache_size": 1000,
|
||||
"max_workers": 8,
|
||||
"enable_batching": True,
|
||||
"batch_size": 10,
|
||||
"max_tokens_per_batch": 2000
|
||||
}
|
||||
self._load_config_file(config_file)
|
||||
self._load_env_vars()
|
||||
|
||||
def get_optimization_config(self) -> Dict:
|
||||
"""Get optimization configuration."""
|
||||
return self._configs.get("optimization", {})
|
||||
|
||||
|
||||
def _load_config_file(self, config_file: Optional[str]):
|
||||
"""Load configuration from file."""
|
||||
if config_file and Path(config_file).exists():
|
||||
@@ -114,6 +129,66 @@ class Config:
|
||||
return self._configs[provider].get("api_key")
|
||||
return os.getenv(f"{provider.upper()}_API_KEY")
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""
|
||||
Get configuration value by key.
|
||||
Searches in top-level configs and optimization settings.
|
||||
"""
|
||||
# 1. Check top-level keys
|
||||
if key in self._configs:
|
||||
return self._configs[key]
|
||||
|
||||
# 2. Check optimization settings (common keys)
|
||||
if "optimization" in self._configs and key in self._configs["optimization"]:
|
||||
return self._configs["optimization"][key]
|
||||
|
||||
# 3. Handle specific mapping for optimization keys
|
||||
# Map cache_enabled -> enable_cache if needed
|
||||
if key == "cache_enabled":
|
||||
return self._configs.get("optimization", {}).get("enable_cache", default)
|
||||
|
||||
return default
|
||||
|
||||
|
||||
# Global config instance
|
||||
config = Config()
|
||||
|
||||
|
||||
def resolve_max_workers(
|
||||
explicit: Optional[int] = None,
|
||||
local_config: Optional[Dict[str, Any]] = None,
|
||||
methods: Optional[Any] = None,
|
||||
) -> int:
|
||||
def to_int(val: Any, default: int) -> int:
|
||||
try:
|
||||
return int(val)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
if isinstance(methods, str):
|
||||
normalized_methods = [methods]
|
||||
elif isinstance(methods, (list, tuple, set)):
|
||||
normalized_methods = [m for m in methods if isinstance(m, str)]
|
||||
else:
|
||||
normalized_methods = []
|
||||
|
||||
if explicit is not None:
|
||||
value = to_int(explicit, 1)
|
||||
elif local_config and "max_workers" in local_config:
|
||||
value = to_int(local_config.get("max_workers", 1), 1)
|
||||
else:
|
||||
value = to_int(config.get("max_workers", 5), 5)
|
||||
|
||||
if "ml" in normalized_methods and explicit is None and not (local_config and "max_workers" in local_config):
|
||||
value = 1
|
||||
|
||||
if value < 1:
|
||||
value = 1
|
||||
|
||||
cpu_count = multiprocessing.cpu_count() or 1
|
||||
if value > cpu_count:
|
||||
value = cpu_count
|
||||
if value > 32:
|
||||
value = 32
|
||||
|
||||
return value
|
||||
|
||||
@@ -86,6 +86,7 @@ class CoreferenceChain:
|
||||
mentions: List[Mention]
|
||||
representative: Mention
|
||||
entity_type: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class CoreferenceResolver:
|
||||
@@ -121,12 +122,18 @@ class CoreferenceResolver:
|
||||
)
|
||||
self.chain_builder = CoreferenceChainBuilder(**self.config.get("chain", {}))
|
||||
|
||||
def resolve_coreferences(self, text: str, **options) -> List[CoreferenceChain]:
|
||||
def resolve_coreferences(
|
||||
self,
|
||||
text: str,
|
||||
entities: Optional[List[Entity]] = None,
|
||||
**options
|
||||
) -> List[CoreferenceChain]:
|
||||
"""
|
||||
Resolve coreferences in text.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
entities: List of entities (optional)
|
||||
**options: Resolution options
|
||||
|
||||
Returns:
|
||||
@@ -139,6 +146,8 @@ class CoreferenceResolver:
|
||||
)
|
||||
|
||||
try:
|
||||
from .ner_extractor import NERExtractor
|
||||
|
||||
total_steps = 4 # Extract mentions, resolve pronouns, detect coreferences, build chains
|
||||
current_step = 0
|
||||
|
||||
@@ -151,8 +160,38 @@ class CoreferenceResolver:
|
||||
total=total_steps,
|
||||
message=f"Extracting mentions... ({current_step}/{total_steps}, remaining: {remaining_steps} steps)"
|
||||
)
|
||||
|
||||
# Extract pronouns
|
||||
mentions = self._extract_mentions(text)
|
||||
|
||||
# Add entities as mentions
|
||||
if entities is None:
|
||||
# Extract entities if not provided
|
||||
ner_config = self.config.get("ner", {})
|
||||
if "ner_method" in self.config:
|
||||
ner_config["method"] = self.config["ner_method"]
|
||||
ner = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation", "chain", "entity", "pronoun"]
|
||||
},
|
||||
)
|
||||
entities = ner.extract_entities(text, **options)
|
||||
|
||||
if entities:
|
||||
for entity in entities:
|
||||
mentions.append(
|
||||
Mention(
|
||||
text=entity.text,
|
||||
start_char=entity.start_char,
|
||||
end_char=entity.end_char,
|
||||
mention_type="entity",
|
||||
metadata={"entity_label": entity.label, "confidence": entity.confidence},
|
||||
)
|
||||
)
|
||||
|
||||
# Step 2: Resolve pronouns
|
||||
current_step += 1
|
||||
remaining_steps = total_steps - current_step
|
||||
@@ -201,20 +240,128 @@ class CoreferenceResolver:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
verbose_mode = options.get("verbose", False) or self.config.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [CoreferenceResolver] ERROR: Resolution failed: {e}", flush=True, file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
raise
|
||||
|
||||
def resolve(self, text: str, **options) -> List[CoreferenceChain]:
|
||||
def resolve(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Union[List[CoreferenceChain], List[List[CoreferenceChain]]]:
|
||||
"""
|
||||
Resolve coreferences in text (alias for resolve_coreferences).
|
||||
Resolve coreferences in text or list of documents.
|
||||
Handles batch processing with progress tracking.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
**options: Resolution options
|
||||
text: Input text or list of documents
|
||||
entities: List of entities or list of list of entities (optional)
|
||||
pipeline_id: Optional pipeline ID for progress tracking
|
||||
**kwargs: Resolution options
|
||||
|
||||
Returns:
|
||||
list: List of coreference chains
|
||||
Union[List[CoreferenceChain], List[List[CoreferenceChain]]]: Resolved coreference chains
|
||||
"""
|
||||
return self.resolve_coreferences(text, **options)
|
||||
if isinstance(text, list):
|
||||
# Handle batch resolution with progress tracking
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="CoreferenceResolver",
|
||||
message=f"Batch resolving coreferences from {len(text)} documents",
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
try:
|
||||
results = []
|
||||
total_items = len(text)
|
||||
total_chains_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch resolution... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
for idx, item in enumerate(text):
|
||||
# Prepare arguments for single item
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
|
||||
doc_entities = None
|
||||
if entities and idx < len(entities):
|
||||
doc_entities = entities[idx]
|
||||
|
||||
# Resolve
|
||||
chains = self.resolve_coreferences(doc_text, entities=doc_entities, **kwargs)
|
||||
|
||||
# Add provenance metadata
|
||||
for chain in chains:
|
||||
# Update chain metadata
|
||||
if chain.metadata is None:
|
||||
chain.metadata = {}
|
||||
chain.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
chain.metadata["document_id"] = item["id"]
|
||||
|
||||
# Update mentions metadata
|
||||
for mention in chain.mentions:
|
||||
if mention.metadata is None:
|
||||
mention.metadata = {}
|
||||
mention.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
mention.metadata["document_id"] = item["id"]
|
||||
|
||||
# Update representative metadata
|
||||
if chain.representative:
|
||||
if chain.representative.metadata is None:
|
||||
chain.representative.metadata = {}
|
||||
chain.representative.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
chain.representative.metadata["document_id"] = item["id"]
|
||||
|
||||
results.append(chains)
|
||||
total_chains_count += len(chains)
|
||||
|
||||
# Update progress
|
||||
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||
remaining = total_items - (idx + 1)
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=idx + 1,
|
||||
total=total_items,
|
||||
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining}) - Resolved {total_chains_count} chains"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch resolution completed. Processed {len(results)} documents, resolved {total_chains_count} chains.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
else:
|
||||
# Single item
|
||||
return self.resolve_coreferences(text, entities=entities, **kwargs)
|
||||
|
||||
def _extract_mentions(self, text: str) -> List[Mention]:
|
||||
"""Extract all mentions from text."""
|
||||
mentions = []
|
||||
@@ -346,15 +493,49 @@ class PronounResolver:
|
||||
if m.mention_type == "entity" or m.mention_type == "nominal"
|
||||
]
|
||||
|
||||
# Simple resolution: find closest preceding entity
|
||||
# Simple resolution: find closest preceding entity with compatible type
|
||||
pronoun_types = {
|
||||
"he": ["PERSON"],
|
||||
"him": ["PERSON"],
|
||||
"his": ["PERSON"],
|
||||
"she": ["PERSON"],
|
||||
"her": ["PERSON"],
|
||||
"it": ["ORG", "GPE", "LOC", "PRODUCT", "EVENT", "FAC", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL"],
|
||||
"its": ["ORG", "GPE", "LOC", "PRODUCT", "EVENT", "FAC", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL"],
|
||||
"they": ["ORG", "GPE", "PERSON", "NORP"], # Can be groups of people or organizations
|
||||
"them": ["ORG", "GPE", "PERSON", "NORP"],
|
||||
"their": ["ORG", "GPE", "PERSON", "NORP"],
|
||||
}
|
||||
|
||||
for pronoun in pronouns:
|
||||
# Find preceding entities
|
||||
preceding = [e for e in entities if e.end_char < pronoun.start_char]
|
||||
|
||||
if preceding:
|
||||
# Take closest
|
||||
antecedent = max(preceding, key=lambda e: e.end_char)
|
||||
pronoun_lower = pronoun.text.lower()
|
||||
compatible_types = pronoun_types.get(pronoun_lower)
|
||||
|
||||
antecedent = None
|
||||
|
||||
if compatible_types:
|
||||
# Filter by type
|
||||
compatible = [
|
||||
e for e in preceding
|
||||
if e.metadata and e.metadata.get("entity_label") in compatible_types
|
||||
]
|
||||
if compatible:
|
||||
# Take closest compatible
|
||||
antecedent = max(compatible, key=lambda e: e.end_char)
|
||||
|
||||
# Fallback to closest if no compatible found or pronoun type unknown
|
||||
if antecedent is None:
|
||||
antecedent = max(preceding, key=lambda e: e.end_char)
|
||||
|
||||
resolutions.append((pronoun.text, antecedent.text))
|
||||
|
||||
# Update pronoun metadata and link to antecedent
|
||||
pronoun.entity_id = antecedent.text
|
||||
pronoun.metadata["antecedent_text"] = antecedent.text
|
||||
|
||||
return resolutions
|
||||
|
||||
@@ -431,32 +612,59 @@ class CoreferenceChainBuilder:
|
||||
list: List of coreference chains
|
||||
"""
|
||||
chains = []
|
||||
processed_indices = set()
|
||||
|
||||
# Simple implementation: group by text similarity
|
||||
processed = set()
|
||||
|
||||
for mention in mentions:
|
||||
if mention.text.lower() in processed:
|
||||
for i, mention in enumerate(mentions):
|
||||
if i in processed_indices:
|
||||
continue
|
||||
|
||||
# Find similar mentions
|
||||
similar = [
|
||||
m
|
||||
for m in mentions
|
||||
if m.text.lower() == mention.text.lower()
|
||||
or self._similar_mentions(mention.text, m.text)
|
||||
]
|
||||
# Start a new group
|
||||
group = [mention]
|
||||
processed_indices.add(i)
|
||||
|
||||
if len(similar) > 1:
|
||||
processed.add(mention.text.lower())
|
||||
# Find related mentions
|
||||
for j, other in enumerate(mentions):
|
||||
if j in processed_indices:
|
||||
continue
|
||||
|
||||
# Representative is first (leftmost) mention
|
||||
representative = min(similar, key=lambda m: m.start_char)
|
||||
is_related = False
|
||||
|
||||
# 1. Text similarity
|
||||
if (
|
||||
other.text.lower() == mention.text.lower()
|
||||
or self._similar_mentions(mention.text, other.text)
|
||||
):
|
||||
is_related = True
|
||||
|
||||
# 2. Pronoun resolution (entity_id matches text or entity_id matches entity_id)
|
||||
elif mention.entity_id and (
|
||||
mention.entity_id == other.text
|
||||
or mention.entity_id == other.entity_id
|
||||
):
|
||||
is_related = True
|
||||
elif other.entity_id and (
|
||||
other.entity_id == mention.text
|
||||
or other.entity_id == mention.entity_id
|
||||
):
|
||||
is_related = True
|
||||
|
||||
if is_related:
|
||||
group.append(other)
|
||||
processed_indices.add(j)
|
||||
|
||||
if len(group) > 1:
|
||||
# Representative is first (leftmost) mention, or prefer entity over pronoun
|
||||
# Prefer entity mention as representative
|
||||
entities = [m for m in group if m.mention_type != "pronoun"]
|
||||
if entities:
|
||||
representative = min(entities, key=lambda m: m.start_char)
|
||||
else:
|
||||
representative = min(group, key=lambda m: m.start_char)
|
||||
|
||||
chain = CoreferenceChain(
|
||||
mentions=similar,
|
||||
mentions=group,
|
||||
representative=representative,
|
||||
entity_type=similar[0].metadata.get("entity_label"),
|
||||
entity_type=representative.metadata.get("entity_label"),
|
||||
)
|
||||
chains.append(chain)
|
||||
|
||||
|
||||
@@ -85,79 +85,215 @@ class Event:
|
||||
class EventDetector:
|
||||
"""Event detection and extraction handler."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
event_types: Optional[List[str]] = None,
|
||||
extract_participants: bool = True,
|
||||
extract_location: bool = True,
|
||||
extract_time: bool = True,
|
||||
method: Union[str, List[str]] = None,
|
||||
config=None,
|
||||
**kwargs
|
||||
):
|
||||
def __init__(self, method: str = "llm", **config):
|
||||
"""
|
||||
Initialize event detector.
|
||||
|
||||
Args:
|
||||
event_types: Specific event types to detect (e.g., ["launch", "acquisition"])
|
||||
extract_participants: Whether to extract event participants
|
||||
extract_location: Whether to extract event locations
|
||||
extract_time: Whether to extract temporal information
|
||||
method: Extraction method(s) for underlying NER/relation extractors.
|
||||
Can be passed to ner_method and relation_method in config.
|
||||
config: Legacy config dict (deprecated, use kwargs)
|
||||
**kwargs: Configuration options:
|
||||
- ner_method: Method for NER extraction (if entities need to be extracted)
|
||||
- relation_method: Method for relation extraction (if relations need to be extracted)
|
||||
- Other options passed to sub-components
|
||||
method: Extraction method ("llm", "pattern")
|
||||
**config: Configuration options
|
||||
"""
|
||||
self.logger = get_logger("event_detector")
|
||||
self.config = config or {}
|
||||
self.config.update(kwargs)
|
||||
self.config = config
|
||||
self.method = method
|
||||
self.progress_tracker = get_progress_tracker()
|
||||
|
||||
# Ensure progress tracker is enabled
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
# Store parameters
|
||||
self.event_types_filter = event_types
|
||||
self.extract_participants = extract_participants
|
||||
self.extract_location = extract_location
|
||||
self.extract_time = extract_time
|
||||
# Initialize components
|
||||
self.event_classifier = EventClassifier(**config)
|
||||
self.temporal_processor = TemporalEventProcessor(**config)
|
||||
|
||||
# Configure extraction options
|
||||
self.extract_participants = config.get("extract_participants", True)
|
||||
self.extract_location = config.get("extract_location", True)
|
||||
self.extract_time = config.get("extract_time", True)
|
||||
self.event_types_filter = config.get("event_types", [])
|
||||
|
||||
# Define event patterns
|
||||
self.event_patterns = {
|
||||
"acquisition": r"\b(acquired|acquisition|buying|bought|merger|merged)\b",
|
||||
"partnership": r"\b(partnered|partnership|collaborate|collaboration)\b",
|
||||
"launch": r"\b(launch|launched|releasing|released|unveil|unveiled)\b",
|
||||
"investment": r"\b(invest|invested|investment|funding|raised)\b",
|
||||
"legal": r"\b(sue|sued|lawsuit|litigation|legal action)\b",
|
||||
}
|
||||
|
||||
# Pre-compile location patterns
|
||||
self.location_patterns = [
|
||||
re.compile(r"in\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)"),
|
||||
re.compile(r"at\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)"),
|
||||
]
|
||||
|
||||
# Pre-compile time patterns
|
||||
self.time_patterns = [
|
||||
re.compile(r"on\s+([A-Z][a-z]+\s+\d{1,2},?\s+\d{4})"),
|
||||
re.compile(r"in\s+(\d{4})"),
|
||||
re.compile(r"(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})"),
|
||||
]
|
||||
|
||||
# Store method for passing to extractors if needed
|
||||
if method is not None:
|
||||
self.config["ner_method"] = method
|
||||
self.config["relation_method"] = method
|
||||
|
||||
self.event_classifier = EventClassifier(**self.config.get("classifier", {}))
|
||||
self.temporal_processor = TemporalEventProcessor(
|
||||
**self.config.get("temporal", {})
|
||||
)
|
||||
self.relationship_extractor = EventRelationshipExtractor(
|
||||
**self.config.get("relationship", {})
|
||||
)
|
||||
def extract(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Union[List[Event], List[List[Event]]]:
|
||||
"""
|
||||
Detect events in text or list of documents.
|
||||
Handles batch processing with progress tracking.
|
||||
|
||||
# Event patterns
|
||||
self.event_patterns = {
|
||||
"founded": r"founded|created|established",
|
||||
"acquired": r"acquired|bought|purchased",
|
||||
"launched": r"launched|released|introduced",
|
||||
"announced": r"announced|declared|stated",
|
||||
"meeting": r"met|meeting|conference|summit",
|
||||
}
|
||||
Args:
|
||||
text: Input text or list of documents
|
||||
pipeline_id: Optional pipeline ID for progress tracking
|
||||
**kwargs: Detection options
|
||||
|
||||
def detect_events(self, text: str, **options) -> List[Event]:
|
||||
Returns:
|
||||
Union[List[Event], List[List[Event]]]: Detected events
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
# Handle batch detection with progress tracking
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="EventDetector",
|
||||
message=f"Batch detecting events from {len(text)} documents",
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
try:
|
||||
results = [None] * len(text) # Pre-allocate to maintain order
|
||||
total_items = len(text)
|
||||
total_events_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch detection... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=[self.config.get("ner_method"), self.config.get("relation_method"), self.config.get("method")],
|
||||
)
|
||||
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
# Prepare arguments for single item
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
|
||||
# Detect
|
||||
events = self.detect_events(doc_text, **kwargs)
|
||||
|
||||
# Add provenance metadata
|
||||
for event in events:
|
||||
if event.metadata is None:
|
||||
event.metadata = {}
|
||||
event.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
event.metadata["document_id"] = item["id"]
|
||||
|
||||
return idx, events
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing item {idx}: {e}")
|
||||
# Return empty list on failure to continue processing
|
||||
return idx, []
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit tasks
|
||||
future_to_idx = {}
|
||||
for idx, item in enumerate(text):
|
||||
future = executor.submit(process_item, idx, item)
|
||||
future_to_idx[future] = idx
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, events = future.result()
|
||||
results[idx] = events
|
||||
total_events_count += len(events)
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
if processed_count % update_interval == 0 or processed_count == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining}) - Detected {total_events_count} events"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
for idx, item in enumerate(text):
|
||||
_, events = process_item(idx, item)
|
||||
results[idx] = events
|
||||
total_events_count += len(events)
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
if processed_count % update_interval == 0 or processed_count == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining}) - Detected {total_events_count} events"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch detection completed. Processed {len(results)} documents, detected {total_events_count} events.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
else:
|
||||
# Single item
|
||||
return self.detect_events(text, **kwargs)
|
||||
|
||||
def detect_events(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Event], List[List[Event]]]:
|
||||
"""
|
||||
Detect events in text content.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Detection options
|
||||
|
||||
Returns:
|
||||
list: List of detected events
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
return self.extract(text, pipeline_id=pipeline_id, **options)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="EventDetector",
|
||||
|
||||
@@ -12,8 +12,8 @@ Supported Methods (for future extensibility):
|
||||
|
||||
Algorithms Used:
|
||||
- Confidence Thresholding: Statistical threshold-based filtering
|
||||
- Duplicate Detection: Set-based and similarity-based deduplication
|
||||
- Consistency Checking: Rule-based and graph-based consistency validation
|
||||
- Duplicate Detection: (Removed - handled by external module)
|
||||
- Consistency Checking: (Removed - handled by external module)
|
||||
- Quality Scoring: Weighted scoring algorithms for extraction quality
|
||||
- Validation Metrics: Precision, recall, F1-score calculations
|
||||
- Boundary Validation: Character position and text boundary checking
|
||||
@@ -22,7 +22,6 @@ Key Features:
|
||||
- Entity validation with confidence checking
|
||||
- Relation validation and consistency checking
|
||||
- Quality scoring and metrics calculation
|
||||
- Duplicate detection
|
||||
- Confidence-based filtering
|
||||
- Validation result reporting
|
||||
- Method parameter support for future method-specific validation
|
||||
@@ -47,8 +46,10 @@ Author: Semantica Contributors
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional, Set, Tuple, Union
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
from datetime import datetime
|
||||
import re
|
||||
|
||||
from ..utils.exceptions import ProcessingError
|
||||
from ..utils.logging import get_logger
|
||||
@@ -66,6 +67,7 @@ class ValidationResult:
|
||||
errors: List[str] = field(default_factory=list)
|
||||
warnings: List[str] = field(default_factory=list)
|
||||
metrics: Dict[str, Any] = field(default_factory=dict)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class ExtractionValidator:
|
||||
@@ -79,7 +81,6 @@ class ExtractionValidator:
|
||||
method: Validation method (for future extensibility, currently unused)
|
||||
**config: Configuration options:
|
||||
- min_confidence: Minimum confidence threshold (default: 0.5)
|
||||
- validate_consistency: Check consistency (default: True)
|
||||
"""
|
||||
self.logger = get_logger("extraction_validator")
|
||||
self.config = config
|
||||
@@ -90,19 +91,30 @@ class ExtractionValidator:
|
||||
|
||||
self.method = method # Reserved for future method-based validation
|
||||
self.min_confidence = config.get("min_confidence", 0.5)
|
||||
self.validate_consistency = config.get("validate_consistency", True)
|
||||
|
||||
def validate_entities(self, entities: List[Entity], **options) -> ValidationResult:
|
||||
def validate_entities(self, entities: Union[List[Entity], List[List[Entity]]], **options) -> Union[ValidationResult, List[ValidationResult]]:
|
||||
"""
|
||||
Validate extracted entities.
|
||||
Handles both single list and batch list of entities.
|
||||
|
||||
Args:
|
||||
entities: List of entities
|
||||
entities: List of entities or list of list of entities
|
||||
**options: Validation options
|
||||
|
||||
Returns:
|
||||
ValidationResult: Validation result
|
||||
ValidationResult or List[ValidationResult]: Validation result(s)
|
||||
"""
|
||||
# Handle batch validation
|
||||
if entities and isinstance(entities, list) and len(entities) > 0 and isinstance(entities[0], list):
|
||||
results = []
|
||||
for idx, batch_entities in enumerate(entities):
|
||||
res = self.validate_entities(batch_entities, **options)
|
||||
# Ensure metadata has batch index
|
||||
if "batch_index" not in res.metadata:
|
||||
res.metadata["batch_index"] = idx
|
||||
results.append(res)
|
||||
return results
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="ExtractionValidator",
|
||||
@@ -126,14 +138,6 @@ class ExtractionValidator:
|
||||
f"{len(low_confidence)} entities below confidence threshold"
|
||||
)
|
||||
|
||||
# Check for duplicates
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Checking for duplicates..."
|
||||
)
|
||||
entity_texts = [e.text.lower() for e in entities]
|
||||
duplicates = len(entity_texts) - len(set(entity_texts))
|
||||
if duplicates > 0:
|
||||
warnings.append(f"{duplicates} duplicate entities found")
|
||||
|
||||
# Check for empty entities
|
||||
empty_entities = [e for e in entities if not e.text.strip()]
|
||||
@@ -148,8 +152,7 @@ class ExtractionValidator:
|
||||
[e for e in entities if min_confidence <= e.confidence < 0.8]
|
||||
),
|
||||
"low_confidence": len(low_confidence),
|
||||
"unique_entities": len(set(entity_texts)),
|
||||
"duplicates": duplicates,
|
||||
"unique_entities": len(set(e.text for e in entities)),
|
||||
"entity_types": len(set(e.label for e in entities)),
|
||||
"average_confidence": sum(e.confidence for e in entities)
|
||||
/ len(entities)
|
||||
@@ -162,12 +165,23 @@ class ExtractionValidator:
|
||||
|
||||
valid = len(errors) == 0
|
||||
|
||||
# Collect metadata from entities
|
||||
metadata = {}
|
||||
if entities:
|
||||
first = entities[0]
|
||||
if hasattr(first, "metadata") and first.metadata:
|
||||
if "batch_index" in first.metadata:
|
||||
metadata["batch_index"] = first.metadata["batch_index"]
|
||||
if "document_id" in first.metadata:
|
||||
metadata["document_id"] = first.metadata["document_id"]
|
||||
|
||||
result = ValidationResult(
|
||||
valid=valid,
|
||||
score=score,
|
||||
errors=errors,
|
||||
warnings=warnings,
|
||||
metrics=metrics,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
@@ -184,18 +198,30 @@ class ExtractionValidator:
|
||||
raise
|
||||
|
||||
def validate_relations(
|
||||
self, relations: List[Relation], **options
|
||||
) -> ValidationResult:
|
||||
self, relations: Union[List[Relation], List[List[Relation]]], **options
|
||||
) -> Union[ValidationResult, List[ValidationResult]]:
|
||||
"""
|
||||
Validate extracted relations.
|
||||
Handles both single list and batch list of relations.
|
||||
|
||||
Args:
|
||||
relations: List of relations
|
||||
relations: List of relations or list of list of relations
|
||||
**options: Validation options
|
||||
|
||||
Returns:
|
||||
ValidationResult: Validation result
|
||||
ValidationResult or List[ValidationResult]: Validation result(s)
|
||||
"""
|
||||
# Handle batch validation
|
||||
if relations and isinstance(relations, list) and len(relations) > 0 and isinstance(relations[0], list):
|
||||
results = []
|
||||
for idx, batch_relations in enumerate(relations):
|
||||
res = self.validate_relations(batch_relations, **options)
|
||||
# Ensure metadata has batch index
|
||||
if "batch_index" not in res.metadata:
|
||||
res.metadata["batch_index"] = idx
|
||||
results.append(res)
|
||||
return results
|
||||
|
||||
errors = []
|
||||
warnings = []
|
||||
metrics = {}
|
||||
@@ -218,12 +244,6 @@ class ExtractionValidator:
|
||||
if invalid_relations:
|
||||
errors.append(f"{len(invalid_relations)} invalid relations found")
|
||||
|
||||
# Check consistency
|
||||
if self.validate_consistency:
|
||||
consistency_issues = self._check_consistency(relations)
|
||||
if consistency_issues:
|
||||
warnings.append(f"{len(consistency_issues)} consistency issues found")
|
||||
|
||||
# Calculate metrics
|
||||
metrics = {
|
||||
"total_relations": len(relations),
|
||||
@@ -244,31 +264,25 @@ class ExtractionValidator:
|
||||
|
||||
valid = len(errors) == 0
|
||||
|
||||
# Collect metadata from relations
|
||||
metadata = {}
|
||||
if relations:
|
||||
first = relations[0]
|
||||
if hasattr(first, "metadata") and first.metadata:
|
||||
if "batch_index" in first.metadata:
|
||||
metadata["batch_index"] = first.metadata["batch_index"]
|
||||
if "document_id" in first.metadata:
|
||||
metadata["document_id"] = first.metadata["document_id"]
|
||||
|
||||
return ValidationResult(
|
||||
valid=valid, score=score, errors=errors, warnings=warnings, metrics=metrics
|
||||
valid=valid,
|
||||
score=score,
|
||||
errors=errors,
|
||||
warnings=warnings,
|
||||
metrics=metrics,
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
def _check_consistency(self, relations: List[Relation]) -> List[str]:
|
||||
"""Check consistency of relations."""
|
||||
issues = []
|
||||
|
||||
# Check for contradictory relations
|
||||
relation_pairs = {}
|
||||
for relation in relations:
|
||||
key = (relation.subject.text, relation.object.text)
|
||||
if key not in relation_pairs:
|
||||
relation_pairs[key] = []
|
||||
relation_pairs[key].append(relation.predicate)
|
||||
|
||||
# Find contradictions (e.g., "founded_by" and "founded" for same pair)
|
||||
for key, predicates in relation_pairs.items():
|
||||
if len(set(predicates)) > 1:
|
||||
# Check for obvious contradictions
|
||||
if "founded_by" in predicates and "founded" in predicates:
|
||||
issues.append(f"Contradictory relations for {key}")
|
||||
|
||||
return issues
|
||||
|
||||
def _calculate_entity_score(
|
||||
self, entities: List[Entity], metrics: Dict[str, Any]
|
||||
) -> float:
|
||||
|
||||
@@ -108,7 +108,12 @@ class LLMExtraction:
|
||||
|
||||
# Initialize provider using new system
|
||||
try:
|
||||
self.provider = create_provider(provider, **config)
|
||||
# Sanitize config: remove api_key if it's None/empty to allow fallback
|
||||
provider_config = config.copy()
|
||||
if "api_key" in provider_config and not provider_config["api_key"]:
|
||||
del provider_config["api_key"]
|
||||
|
||||
self.provider = create_provider(provider, **provider_config)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to initialize {provider} provider: {e}")
|
||||
self.provider = None
|
||||
@@ -289,7 +294,14 @@ Return the enhanced relation list in JSON format."""
|
||||
) -> List[Entity]:
|
||||
"""Parse LLM response for entities."""
|
||||
# Simplified parsing - in practice would parse JSON
|
||||
# For now, return original entities
|
||||
# For now, return original entities with updated metadata
|
||||
for entity in original_entities:
|
||||
if entity.metadata is None:
|
||||
entity.metadata = {}
|
||||
entity.metadata.update({
|
||||
"enhanced_by": self.provider_name,
|
||||
"model": self.model
|
||||
})
|
||||
return original_entities
|
||||
|
||||
def _parse_relation_response(
|
||||
@@ -297,7 +309,14 @@ Return the enhanced relation list in JSON format."""
|
||||
) -> List[Relation]:
|
||||
"""Parse LLM response for relations."""
|
||||
# Simplified parsing - in practice would parse JSON
|
||||
# For now, return original relations
|
||||
# For now, return original relations with updated metadata
|
||||
for relation in original_relations:
|
||||
if relation.metadata is None:
|
||||
relation.metadata = {}
|
||||
relation.metadata.update({
|
||||
"enhanced_by": self.provider_name,
|
||||
"model": self.model
|
||||
})
|
||||
return original_relations
|
||||
|
||||
|
||||
|
||||
+1049
-194
File diff suppressed because it is too large
Load Diff
@@ -20,7 +20,12 @@ Algorithms Used:
|
||||
- Transformer Models: BERT, RoBERTa, DistilBERT for token classification
|
||||
- Large Language Models: GPT, Claude, Gemini for zero-shot/few-shot extraction
|
||||
- Ensemble Voting: Majority voting and confidence-weighted aggregation
|
||||
- Deduplication: Set-based and similarity-based entity deduplication
|
||||
- Weighted Confidence Scoring:
|
||||
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||
* Method_Confidence: Confidence score from the extraction algorithm
|
||||
* Type_Similarity_Score: Semantic match with user-provided entity types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||
- Last Resort Fallback: Capitalized word heuristic when all other methods fail
|
||||
|
||||
Key Features:
|
||||
- Multiple extraction methods:
|
||||
@@ -31,8 +36,9 @@ Key Features:
|
||||
* HuggingFace: Custom HuggingFace NER models
|
||||
* LLM-based: Large language model extraction
|
||||
- Fallback chain support: Try methods in order until one succeeds
|
||||
- Robust Fallbacks: Prevents empty results via ML -> Pattern -> Last Resort chain
|
||||
- Ensemble voting: Combine results from multiple methods
|
||||
- Post-processing: Entity boundary validation and deduplication
|
||||
- Post-processing: Entity boundary validation
|
||||
- Multiple entity type support (PERSON, ORG, GPE, DATE, etc.)
|
||||
- Confidence scoring and filtering
|
||||
- Batch processing capabilities
|
||||
@@ -90,7 +96,12 @@ class Entity:
|
||||
class NERExtractor:
|
||||
"""Named Entity Recognition extractor."""
|
||||
|
||||
def __init__(self, method: Union[str, List[str]] = "ml", **config):
|
||||
def __init__(
|
||||
self,
|
||||
method: Union[str, List[str]] = "ml",
|
||||
entity_types: Optional[List[str]] = None,
|
||||
**config
|
||||
):
|
||||
"""
|
||||
Initialize NER extractor.
|
||||
|
||||
@@ -103,6 +114,8 @@ class NERExtractor:
|
||||
- "huggingface": HuggingFace model
|
||||
- "llm": LLM-based extraction
|
||||
- List of methods for fallback chain
|
||||
entity_types: List of entity types to extract (e.g., ["PERSON", "ORG"]).
|
||||
If provided, extraction methods will try to limit/focus on these types.
|
||||
**config: Configuration options:
|
||||
- model: Model name (for ML/HuggingFace methods)
|
||||
- huggingface_model: HuggingFace model name
|
||||
@@ -115,6 +128,7 @@ class NERExtractor:
|
||||
"""
|
||||
self.logger = get_logger("ner_extractor")
|
||||
self.config = config
|
||||
self.entity_types = entity_types
|
||||
|
||||
# Method configuration
|
||||
self.method = method if isinstance(method, list) else [method]
|
||||
@@ -164,8 +178,11 @@ class NERExtractor:
|
||||
)
|
||||
|
||||
try:
|
||||
results = []
|
||||
results = [None] * len(text)
|
||||
total_items = len(text)
|
||||
total_entities_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
||||
if total_items <= 10:
|
||||
update_interval = 1 # Update every item for small datasets
|
||||
@@ -173,49 +190,107 @@ class NERExtractor:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update - ALWAYS show this
|
||||
remaining = total_items
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch extraction... 0/{total_items} (remaining: {remaining})"
|
||||
message=f"Starting batch extraction... 0/{total_items}"
|
||||
)
|
||||
|
||||
for idx, item in enumerate(text, 1):
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
# Helper function for single item processing
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
current_entities = []
|
||||
if isinstance(item, dict) and "content" in item:
|
||||
results.append(self.extract_entities(item["content"], **kwargs))
|
||||
current_entities = self.extract_entities(item["content"], **kwargs)
|
||||
elif isinstance(item, str):
|
||||
results.append(self.extract_entities(item, **kwargs))
|
||||
current_entities = self.extract_entities(item, **kwargs)
|
||||
else:
|
||||
# Try converting to string
|
||||
try:
|
||||
results.append(self.extract_entities(str(item), **kwargs))
|
||||
current_entities = self.extract_entities(str(item), **kwargs)
|
||||
except Exception:
|
||||
results.append([])
|
||||
except Exception:
|
||||
results.append([])
|
||||
current_entities = []
|
||||
|
||||
# Add provenance metadata
|
||||
for ent in current_entities:
|
||||
if ent.metadata is None:
|
||||
ent.metadata = {}
|
||||
ent.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
ent.metadata["document_id"] = item["id"]
|
||||
|
||||
return idx, current_entities
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to process item {idx}: {e}")
|
||||
return idx, []
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
remaining = total_items - idx
|
||||
# Update progress: always update for small datasets, or at intervals for large ones
|
||||
should_update = (
|
||||
idx % update_interval == 0 or
|
||||
idx == total_items or
|
||||
idx == 1 or
|
||||
total_items <= 10 # Always update for small datasets
|
||||
)
|
||||
if should_update:
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=idx,
|
||||
total=total_items,
|
||||
message=f"Processing documents... {idx}/{total_items} (remaining: {remaining})"
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit all tasks
|
||||
future_to_idx = {
|
||||
executor.submit(process_item, idx, item): idx
|
||||
for idx, item in enumerate(text)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, entities = future.result()
|
||||
results[idx] = entities
|
||||
total_entities_count += len(entities)
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
should_update = (
|
||||
processed_count % update_interval == 0 or
|
||||
processed_count == total_items or
|
||||
processed_count == 1 or
|
||||
total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing documents... {processed_count}/{total_items} (remaining: {remaining}) - Extracted {total_entities_count} entities so far"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
for idx, item in enumerate(text):
|
||||
_, entities = process_item(idx, item)
|
||||
results[idx] = entities
|
||||
total_entities_count += len(entities)
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
should_update = (
|
||||
processed_count % update_interval == 0 or
|
||||
processed_count == total_items or
|
||||
processed_count == 1 or
|
||||
total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing documents... {processed_count}/{total_items} (remaining: {remaining}) - Extracted {total_entities_count} entities so far"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Extracted entities from {len(results)} documents",
|
||||
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_entities_count} entities.",
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
@@ -226,12 +301,18 @@ class NERExtractor:
|
||||
else:
|
||||
return self.extract_entities(text, **kwargs)
|
||||
|
||||
def extract_entities(self, text: str, **options) -> List[Entity]:
|
||||
def extract_entities(
|
||||
self,
|
||||
text: Union[str, List[Dict[str, Any]], List[str]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Entity], List[List[Entity]]]:
|
||||
"""
|
||||
Extract named entities from text.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options:
|
||||
- entity_types: Filter by entity types (list)
|
||||
- min_confidence: Minimum confidence threshold
|
||||
@@ -240,6 +321,9 @@ class NERExtractor:
|
||||
Returns:
|
||||
list: List of extracted entities
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
return self.extract(text, pipeline_id=pipeline_id, **options)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="NERExtractor",
|
||||
@@ -260,10 +344,12 @@ class NERExtractor:
|
||||
methods = [methods]
|
||||
|
||||
min_confidence = options.get("min_confidence", self.min_confidence)
|
||||
entity_types = options.get("entity_types")
|
||||
entity_types = options.get("entity_types", self.entity_types)
|
||||
|
||||
# Merge config with options
|
||||
all_options = {**self.config, **options}
|
||||
if entity_types:
|
||||
all_options["entity_types"] = entity_types
|
||||
|
||||
# Try each method in order (fallback chain)
|
||||
all_entities = []
|
||||
@@ -289,40 +375,46 @@ class NERExtractor:
|
||||
method_options["model"] = all_options.get(
|
||||
"llm_model", all_options.get("model")
|
||||
)
|
||||
# Pass api_key if provided (needed for all providers)
|
||||
if "api_key" in all_options:
|
||||
method_options["api_key"] = all_options["api_key"]
|
||||
elif "api_key" not in method_options:
|
||||
# Try to get from environment as fallback
|
||||
# Ensure api_key is populated: check explicitly provided or fallback to env
|
||||
current_key = method_options.get("api_key")
|
||||
if not current_key:
|
||||
# Not found or empty/None, try environment
|
||||
import os
|
||||
provider = method_options.get("provider", "openai")
|
||||
env_key = f"{provider.upper()}_API_KEY"
|
||||
provider_name = method_options.get("provider", "openai")
|
||||
env_key = f"{provider_name.upper()}_API_KEY"
|
||||
api_key = os.getenv(env_key)
|
||||
if api_key:
|
||||
method_options["api_key"] = api_key
|
||||
# Pass entity_types to LLM method so it can use them in the prompt
|
||||
if entity_types:
|
||||
method_options["entity_types"] = entity_types
|
||||
|
||||
entities = method_func(text, **method_options)
|
||||
|
||||
# Filter by confidence and entity types
|
||||
filtered = [e for e in entities if e.confidence >= min_confidence]
|
||||
# Apply weighted scoring if entity_types are provided
|
||||
if entity_types:
|
||||
# Case-insensitive and flexible matching for entity types
|
||||
entity_types_lower = {et.lower() for et in entity_types}
|
||||
filtered = [
|
||||
e for e in filtered
|
||||
if e.label.lower() in entity_types_lower
|
||||
or any(et.lower() in e.label.lower() or e.label.lower() in et.lower()
|
||||
for et in entity_types)
|
||||
]
|
||||
try:
|
||||
from .methods import calculate_weighted_confidence
|
||||
for e in entities:
|
||||
e.confidence = calculate_weighted_confidence(
|
||||
item_type=e.label,
|
||||
original_confidence=e.confidence,
|
||||
valid_types=entity_types,
|
||||
item_text=e.text
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Filter by confidence
|
||||
filtered = [e for e in entities if e.confidence >= min_confidence]
|
||||
|
||||
if filtered:
|
||||
all_entities.append((method_name, filtered))
|
||||
|
||||
# If not using ensemble, return first successful result
|
||||
if not self.ensemble_voting:
|
||||
# Ensure default metadata
|
||||
for e in filtered:
|
||||
if e.metadata is None: e.metadata = {}
|
||||
if "batch_index" not in e.metadata: e.metadata["batch_index"] = 0
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
@@ -342,7 +434,8 @@ class NERExtractor:
|
||||
elif all_entities:
|
||||
entities = all_entities[0][1] # Use first successful method
|
||||
else:
|
||||
entities = []
|
||||
# Fallback to pattern-based extraction if all models fail
|
||||
entities = self._extract_fallback(text)
|
||||
|
||||
# Post-processing if enabled
|
||||
if self.post_process and entities:
|
||||
@@ -390,19 +483,12 @@ class NERExtractor:
|
||||
def _post_process_entities(self, entities: List[Entity], text: str) -> List[Entity]:
|
||||
"""Post-process entities for refinement."""
|
||||
processed = []
|
||||
seen = set()
|
||||
|
||||
for entity in entities:
|
||||
# Check boundaries
|
||||
if entity.start_char < 0 or entity.end_char > len(text):
|
||||
continue
|
||||
|
||||
# Check for duplicates
|
||||
key = (entity.text.lower(), entity.label, entity.start_char)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
|
||||
# Validate entity text matches
|
||||
actual_text = text[entity.start_char : entity.end_char]
|
||||
if actual_text.lower() != entity.text.lower():
|
||||
@@ -457,6 +543,7 @@ class NERExtractor:
|
||||
def _extract_fallback(self, text: str) -> List[Entity]:
|
||||
"""Fallback entity extraction using simple patterns."""
|
||||
entities = []
|
||||
import re
|
||||
|
||||
# Simple patterns for common entity types
|
||||
patterns = {
|
||||
@@ -466,20 +553,49 @@ class NERExtractor:
|
||||
"DATE": r"\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4})\b",
|
||||
}
|
||||
|
||||
import re
|
||||
# Track covered ranges to avoid overlaps
|
||||
covered_ranges = set()
|
||||
|
||||
for label, pattern in patterns.items():
|
||||
for match in re.finditer(pattern, text):
|
||||
entities.append(
|
||||
Entity(
|
||||
text=match.group(1),
|
||||
label=label,
|
||||
start_char=match.start(),
|
||||
end_char=match.end(),
|
||||
confidence=0.7, # Lower confidence for pattern-based
|
||||
metadata={"extraction_method": "pattern"},
|
||||
start, end = match.start(), match.end()
|
||||
# Check overlap
|
||||
is_overlap = any(r_start < end and r_end > start for r_start, r_end in covered_ranges)
|
||||
if not is_overlap:
|
||||
# Use group 1 if available, else group 0
|
||||
text_val = match.group(1) if match.lastindex and match.lastindex >= 1 else match.group(0)
|
||||
|
||||
entities.append(
|
||||
Entity(
|
||||
text=text_val,
|
||||
label=label,
|
||||
start_char=start,
|
||||
end_char=end,
|
||||
confidence=0.7, # Lower confidence for pattern-based
|
||||
metadata={"extraction_method": "pattern"},
|
||||
)
|
||||
)
|
||||
)
|
||||
covered_ranges.add((start, end))
|
||||
|
||||
# Last Resort: If no entities found, try single capitalized words as generic entities
|
||||
if not entities:
|
||||
# Match any capitalized word of length > 2
|
||||
cap_pattern = r"\b[A-Z][a-z]{2,}\b"
|
||||
for match in re.finditer(cap_pattern, text):
|
||||
start, end = match.start(), match.end()
|
||||
is_overlap = any(r_start < end and r_end > start for r_start, r_end in covered_ranges)
|
||||
if not is_overlap:
|
||||
entities.append(
|
||||
Entity(
|
||||
text=match.group(0),
|
||||
label="UNKNOWN",
|
||||
start_char=start,
|
||||
end_char=end,
|
||||
confidence=0.5,
|
||||
metadata={"extraction_method": "last_resort_pattern"},
|
||||
)
|
||||
)
|
||||
covered_ranges.add((start, end))
|
||||
|
||||
return entities
|
||||
|
||||
@@ -494,7 +610,7 @@ class NERExtractor:
|
||||
Returns:
|
||||
list: List of entity lists for each text
|
||||
"""
|
||||
return [self.extract_entities(text, **options) for text in texts]
|
||||
return self.extract(texts, **options)
|
||||
|
||||
def classify_entities(self, entities: List[Entity]) -> Dict[str, List[Entity]]:
|
||||
"""
|
||||
|
||||
@@ -71,7 +71,19 @@ License: MIT
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Union, Type
|
||||
|
||||
try:
|
||||
from pydantic import BaseModel, ValidationError
|
||||
except ImportError:
|
||||
BaseModel = Any
|
||||
ValidationError = Exception
|
||||
|
||||
try:
|
||||
import instructor
|
||||
except ImportError:
|
||||
instructor = None
|
||||
|
||||
from ..utils.exceptions import ProcessingError
|
||||
from ..utils.logging import get_logger
|
||||
@@ -211,6 +223,291 @@ class BaseProvider:
|
||||
raise ProcessingError(f"Failed to generate structured output: {last_error}")
|
||||
return []
|
||||
|
||||
def generate_typed(
|
||||
self,
|
||||
prompt: str,
|
||||
schema: Type[BaseModel],
|
||||
max_retries: int = 3,
|
||||
**kwargs
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate structured output validated against a Pydantic schema.
|
||||
Uses instructor if available and supported for the provider, otherwise falls back to a repair loop.
|
||||
"""
|
||||
provider_name = self.__class__.__name__
|
||||
|
||||
# Try using instructor first if available
|
||||
if instructor:
|
||||
try:
|
||||
client = None
|
||||
mode = instructor.Mode.TOOLS # Default mode
|
||||
|
||||
if provider_name == "OpenAIProvider" and self.client:
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"openai/{kwargs.get('model', self.model)}",
|
||||
api_key=self.api_key
|
||||
)
|
||||
except Exception:
|
||||
client = instructor.from_openai(self.client)
|
||||
else:
|
||||
client = instructor.from_openai(self.client)
|
||||
elif provider_name == "AnthropicProvider" and self.client:
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"anthropic/{kwargs.get('model', self.model)}",
|
||||
api_key=self.api_key
|
||||
)
|
||||
except Exception:
|
||||
client = instructor.from_anthropic(self.client)
|
||||
else:
|
||||
client = instructor.from_anthropic(self.client)
|
||||
elif provider_name == "GeminiProvider" and self.client:
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"gemini/{kwargs.get('model', self.model)}",
|
||||
api_key=self.api_key
|
||||
)
|
||||
except Exception:
|
||||
client = instructor.from_gemini(
|
||||
self.client,
|
||||
mode=instructor.Mode.GEMINI_JSON
|
||||
)
|
||||
else:
|
||||
client = instructor.from_gemini(
|
||||
self.client,
|
||||
mode=instructor.Mode.GEMINI_JSON
|
||||
)
|
||||
elif provider_name == "GroqProvider" and self.client:
|
||||
# Try using from_provider which is recommended for Groq in latest instructor
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"groq/{kwargs.get('model', self.model)}",
|
||||
api_key=self.api_key
|
||||
)
|
||||
except Exception:
|
||||
client = None
|
||||
|
||||
if not client:
|
||||
# Try using from_groq if available (newer instructor versions)
|
||||
if hasattr(instructor, "from_groq"):
|
||||
client = instructor.from_groq(self.client, mode=instructor.Mode.JSON)
|
||||
else:
|
||||
# Fallback: Create OpenAI client pointing to Groq
|
||||
# This avoids the "Client should be an instance of openai.OpenAI" warning
|
||||
try:
|
||||
from openai import OpenAI
|
||||
# Fix: Use self.api_key instead of self.client.api_key
|
||||
groq_client = OpenAI(
|
||||
base_url="https://api.groq.com/openai/v1",
|
||||
api_key=self.api_key,
|
||||
)
|
||||
client = instructor.from_openai(groq_client, mode=instructor.Mode.JSON)
|
||||
except Exception:
|
||||
# Last resort: try passing the groq client directly
|
||||
client = instructor.from_openai(self.client, mode=instructor.Mode.JSON)
|
||||
elif provider_name == "OllamaProvider":
|
||||
# Try from_provider for Ollama if available
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"ollama/{kwargs.get('model', self.model)}",
|
||||
)
|
||||
except Exception:
|
||||
client = None
|
||||
|
||||
if not client:
|
||||
# Create OpenAI-compatible client for Ollama
|
||||
try:
|
||||
from openai import OpenAI
|
||||
# Ollama typically runs on localhost:11434/v1
|
||||
base_url = getattr(self, "base_url", "http://localhost:11434")
|
||||
if not base_url.endswith("/v1"):
|
||||
base_url = f"{base_url.rstrip('/')}/v1"
|
||||
|
||||
ollama_client = OpenAI(
|
||||
base_url=base_url,
|
||||
api_key="ollama", # required but unused
|
||||
)
|
||||
client = instructor.from_openai(ollama_client, mode=instructor.Mode.JSON)
|
||||
except ImportError:
|
||||
pass
|
||||
elif provider_name == "DeepSeekProvider" and self.client:
|
||||
# Try from_provider for DeepSeek
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
client = instructor.from_provider(
|
||||
provider=f"deepseek/{kwargs.get('model', self.model)}",
|
||||
api_key=self.api_key
|
||||
)
|
||||
except Exception:
|
||||
client = None
|
||||
|
||||
if not client:
|
||||
# DeepSeek is OpenAI compatible
|
||||
try:
|
||||
from openai import OpenAI
|
||||
if isinstance(self.client, OpenAI):
|
||||
client = instructor.from_openai(self.client, mode=instructor.Mode.JSON)
|
||||
else:
|
||||
# Try creating fresh client
|
||||
ds_client = OpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url="https://api.deepseek.com"
|
||||
)
|
||||
client = instructor.from_openai(ds_client, mode=instructor.Mode.JSON)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Global LiteLLM support - if litellm is passed in kwargs or config
|
||||
if not client and (kwargs.get("litellm") or self.config.get("litellm")):
|
||||
if hasattr(instructor, "from_provider"):
|
||||
try:
|
||||
# Format for litellm in instructor is litellm/model_name
|
||||
provider_model = kwargs.get("model", self.model)
|
||||
litellm_provider = f"litellm/{provider_model}"
|
||||
client = instructor.from_provider(litellm_provider, api_key=self.api_key)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if client:
|
||||
# Map generate arguments to client arguments
|
||||
# Instructor standardizes on chat.completions.create for OpenAI/Groq/Anthropic/Gemini
|
||||
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response_model": schema,
|
||||
"max_retries": max_retries,
|
||||
"temperature": kwargs.get("temperature", 0.1), # Low temp for structured
|
||||
}
|
||||
|
||||
verbose_mode = kwargs.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [BaseProvider.generate_typed] Using instructor via {provider_name}. Client: {type(client)}", flush=True, file=sys.stdout)
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["max_tokens", "max_completion_tokens", "top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "logit_bias", "user", "top_k"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
# Add provider-specific params
|
||||
if provider_name == "GroqProvider":
|
||||
create_kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
response = client.chat.completions.create(**create_kwargs)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [BaseProvider.generate_typed] Typed response received via instructor ({provider_name}).", flush=True, file=sys.stdout)
|
||||
return response
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Instructor generation failed ({e}), falling back to manual repair loop.")
|
||||
|
||||
# Fallback: Manual repair loop
|
||||
last_error = None
|
||||
current_prompt = prompt
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
# 1. Generate JSON
|
||||
# We use generate_structured to get the dict/list
|
||||
json_result = self.generate_structured(current_prompt, max_retries=1, **kwargs)
|
||||
|
||||
# 2. Validate with Schema
|
||||
# If the result is a list and schema expects a wrapper, or vice versa, we might need adjustment
|
||||
# But we assume the prompt asks for the correct structure matching the schema.
|
||||
|
||||
# Special handling if schema is a wrapper but result is a list
|
||||
if isinstance(json_result, list) and hasattr(schema, "entities") and "entities" in schema.model_fields:
|
||||
# Auto-wrap for entities
|
||||
json_result = {"entities": json_result}
|
||||
|
||||
# Handle categorized dictionary input (e.g. {"PERSON": ["Name"], "ORG": ["Corp"]})
|
||||
elif isinstance(json_result, dict) and hasattr(schema, "entities") and "entities" in schema.model_fields:
|
||||
# Check if it's NOT already in the correct format (i.e., missing "entities" key)
|
||||
if "entities" not in json_result:
|
||||
# Check if values are lists, suggesting categorized output
|
||||
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||
if is_categorized:
|
||||
flat_entities = []
|
||||
for label, items in json_result.items():
|
||||
if isinstance(items, list):
|
||||
for item in items:
|
||||
if isinstance(item, str):
|
||||
flat_entities.append({"text": item, "label": label})
|
||||
elif isinstance(item, dict):
|
||||
# If it's already a dict but nested under label
|
||||
item["label"] = label
|
||||
flat_entities.append(item)
|
||||
json_result = {"entities": flat_entities}
|
||||
|
||||
# Handle categorized dictionary input for relations (e.g. {"founded_by": [{"subject":..., "object":...}]})
|
||||
elif isinstance(json_result, dict) and hasattr(schema, "relations") and "relations" in schema.model_fields:
|
||||
if "relations" not in json_result:
|
||||
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||
if is_categorized:
|
||||
flat_relations = []
|
||||
for label, items in json_result.items():
|
||||
if isinstance(items, list):
|
||||
for item in items:
|
||||
if isinstance(item, dict):
|
||||
# If predicate is missing, use the key as predicate
|
||||
if "predicate" not in item:
|
||||
item["predicate"] = label
|
||||
flat_relations.append(item)
|
||||
json_result = {"relations": flat_relations}
|
||||
|
||||
# Handle categorized dictionary input for triplets
|
||||
elif isinstance(json_result, dict) and hasattr(schema, "triplets") and "triplets" in schema.model_fields:
|
||||
if "triplets" not in json_result:
|
||||
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||
if is_categorized:
|
||||
flat_triplets = []
|
||||
for label, items in json_result.items():
|
||||
if isinstance(items, list):
|
||||
for item in items:
|
||||
if isinstance(item, dict):
|
||||
flat_triplets.append(item)
|
||||
json_result = {"triplets": flat_triplets}
|
||||
|
||||
elif isinstance(json_result, list) and hasattr(schema, "relations") and "relations" in schema.model_fields:
|
||||
json_result = {"relations": json_result}
|
||||
elif isinstance(json_result, list) and hasattr(schema, "triplets") and "triplets" in schema.model_fields:
|
||||
json_result = {"triplets": json_result}
|
||||
|
||||
validated = schema.model_validate(json_result)
|
||||
return validated
|
||||
|
||||
except ValidationError as e:
|
||||
last_error = e
|
||||
error_summary = str(e)
|
||||
# Simplify error summary for the LLM
|
||||
# (You could parse e.errors() for a better message)
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = (attempt + 1) * 1
|
||||
self.logger.warning(f"Schema validation failed (attempt {attempt + 1}): {e}. Retrying with error feedback...")
|
||||
|
||||
# Update prompt with error info
|
||||
current_prompt = f"{prompt}\n\nPrevious response was invalid JSON or didn't match schema:\n{error_summary}\n\nPlease fix the errors and return valid JSON matching the schema."
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
self.logger.error(f"Typed generation failed validation: {e}")
|
||||
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(1)
|
||||
else:
|
||||
self.logger.error(f"Typed generation failed: {e}")
|
||||
|
||||
raise ProcessingError(f"Failed to generate typed output after {max_retries} attempts: {last_error}")
|
||||
|
||||
class OpenAIProvider(BaseProvider):
|
||||
"""OpenAI provider implementation."""
|
||||
|
||||
@@ -248,11 +545,24 @@ class OpenAIProvider(BaseProvider):
|
||||
"OpenAI client not initialized. Set OPENAI_API_KEY or pass api_key."
|
||||
)
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=kwargs.get("temperature", 0.3),
|
||||
)
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": kwargs.get("temperature", 0.3),
|
||||
}
|
||||
|
||||
# Support max_tokens and max_completion_tokens (for o1 models)
|
||||
if "max_completion_tokens" in kwargs:
|
||||
create_kwargs["max_completion_tokens"] = kwargs["max_completion_tokens"]
|
||||
elif "max_tokens" in kwargs:
|
||||
create_kwargs["max_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "logit_bias", "user"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
response = self.client.chat.completions.create(**create_kwargs)
|
||||
return response.choices[0].message.content
|
||||
|
||||
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||
@@ -260,12 +570,25 @@ class OpenAIProvider(BaseProvider):
|
||||
if not self.client:
|
||||
raise ProcessingError("OpenAI client not initialized.")
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
response_format={"type": "json_object"},
|
||||
temperature=kwargs.get("temperature", 0.3),
|
||||
)
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response_format": {"type": "json_object"},
|
||||
"temperature": kwargs.get("temperature", 0.3),
|
||||
}
|
||||
|
||||
# Support max_tokens and max_completion_tokens
|
||||
if "max_completion_tokens" in kwargs:
|
||||
create_kwargs["max_completion_tokens"] = kwargs["max_completion_tokens"]
|
||||
elif "max_tokens" in kwargs:
|
||||
create_kwargs["max_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "logit_bias", "user"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
response = self.client.chat.completions.create(**create_kwargs)
|
||||
try:
|
||||
return self._parse_json(response.choices[0].message.content)
|
||||
except Exception as e:
|
||||
@@ -283,26 +606,40 @@ class GeminiProvider(BaseProvider):
|
||||
self.api_key = api_key or config.get_api_key("gemini")
|
||||
self.model = model
|
||||
self.client = None
|
||||
self._use_new_genai = False
|
||||
self._init_client()
|
||||
|
||||
def _init_client(self):
|
||||
"""Initialize Gemini client."""
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
|
||||
from google import genai as new_genai
|
||||
if self.api_key:
|
||||
genai.configure(api_key=self.api_key)
|
||||
self.client = genai.GenerativeModel(self.model)
|
||||
except (ImportError, OSError):
|
||||
self.client = new_genai.Client(api_key=self.api_key)
|
||||
self._use_new_genai = True
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import google.generativeai as old_genai
|
||||
if self.api_key:
|
||||
old_genai.configure(api_key=self.api_key)
|
||||
self.client = old_genai.GenerativeModel(self.model)
|
||||
self._use_new_genai = False
|
||||
except Exception:
|
||||
self.client = None
|
||||
self.logger.warning(
|
||||
"google-generativeai library not installed. Install with: pip install semantica[llm-gemini]"
|
||||
)
|
||||
self.logger.warning("Gemini SDK not installed. Install with: pip install semantica[llm-gemini]")
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""Check if provider is available."""
|
||||
return self.client is not None
|
||||
|
||||
def _resp_text(self, resp: Any) -> str:
|
||||
if hasattr(resp, "text"):
|
||||
return getattr(resp, "text")
|
||||
try:
|
||||
return resp.candidates[0].content.parts[0].text
|
||||
except Exception:
|
||||
return str(resp)
|
||||
|
||||
def generate(self, prompt: str, **kwargs) -> str:
|
||||
"""Generate text from prompt."""
|
||||
if not self.client:
|
||||
@@ -310,30 +647,53 @@ class GeminiProvider(BaseProvider):
|
||||
"Gemini client not initialized. Set GEMINI_API_KEY or pass api_key."
|
||||
)
|
||||
|
||||
response = self.client.generate_content(
|
||||
prompt, generation_config={"temperature": kwargs.get("temperature", 0.3)}
|
||||
)
|
||||
return response.text
|
||||
if self._use_new_genai:
|
||||
model = kwargs.get("model", self.model)
|
||||
temperature = kwargs.get("temperature", 0.3)
|
||||
create_kwargs = {"model": model, "contents": prompt, "config": {"temperature": temperature}}
|
||||
if "max_tokens" in kwargs:
|
||||
create_kwargs["config"]["max_output_tokens"] = kwargs["max_tokens"]
|
||||
for p in ["top_p", "top_k", "stop_sequences", "candidate_count"]:
|
||||
if p in kwargs:
|
||||
create_kwargs["config"][p] = kwargs[p]
|
||||
resp = self.client.models.generate_content(**create_kwargs)
|
||||
return self._resp_text(resp)
|
||||
else:
|
||||
generation_config = {"temperature": kwargs.get("temperature", 0.3)}
|
||||
if "max_tokens" in kwargs:
|
||||
generation_config["max_output_tokens"] = kwargs["max_tokens"]
|
||||
for param in ["top_p", "top_k", "stop_sequences", "candidate_count"]:
|
||||
if param in kwargs:
|
||||
generation_config[param] = kwargs[param]
|
||||
response = self.client.generate_content(prompt, generation_config=generation_config)
|
||||
return self._resp_text(response)
|
||||
|
||||
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||
"""Generate structured output."""
|
||||
if not self.client:
|
||||
raise ProcessingError("Gemini client not initialized.")
|
||||
|
||||
# Add JSON format instruction to prompt
|
||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||
response = self.client.generate_content(json_prompt)
|
||||
try:
|
||||
return self._parse_json(response.text)
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to parse JSON from Gemini response: {e}")
|
||||
if self._use_new_genai:
|
||||
model = kwargs.get("model", self.model)
|
||||
resp = self.client.models.generate_content(model=model, contents=json_prompt)
|
||||
try:
|
||||
return self._parse_json(self._resp_text(resp))
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to parse JSON from Gemini response: {e}")
|
||||
else:
|
||||
response = self.client.generate_content(json_prompt)
|
||||
try:
|
||||
return self._parse_json(self._resp_text(response))
|
||||
except Exception as e:
|
||||
raise ProcessingError(f"Failed to parse JSON from Gemini response: {e}")
|
||||
|
||||
|
||||
class GroqProvider(BaseProvider):
|
||||
"""Groq provider implementation."""
|
||||
|
||||
def __init__(
|
||||
self, api_key: Optional[str] = None, model: str = "llama2-70b-4096", **kwargs
|
||||
self, api_key: Optional[str] = None, model: str = "llama-3.3-70b-versatile", **kwargs
|
||||
):
|
||||
"""Initialize Groq provider."""
|
||||
super().__init__(**kwargs)
|
||||
@@ -405,11 +765,32 @@ class GroqProvider(BaseProvider):
|
||||
"Groq client not initialized. Set GROQ_API_KEY or pass api_key."
|
||||
)
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=kwargs.get("temperature", 0.3),
|
||||
)
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": kwargs.get("temperature", 0.3),
|
||||
}
|
||||
|
||||
# Support max_tokens and max_completion_tokens
|
||||
if "max_completion_tokens" in kwargs:
|
||||
create_kwargs["max_completion_tokens"] = kwargs["max_completion_tokens"]
|
||||
elif "max_tokens" in kwargs:
|
||||
create_kwargs["max_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "user"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
verbose_mode = kwargs.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [GroqProvider.generate] Sending request to Groq API (model: {create_kwargs['model']})...", flush=True, file=sys.stdout)
|
||||
|
||||
response = self.client.chat.completions.create(**create_kwargs)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [GroqProvider.generate] Response received from Groq.", flush=True, file=sys.stdout)
|
||||
return response.choices[0].message.content
|
||||
|
||||
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||
@@ -417,12 +798,38 @@ class GroqProvider(BaseProvider):
|
||||
if not self.client:
|
||||
raise ProcessingError("Groq client not initialized.")
|
||||
|
||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||
response = self.client.chat.completions.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
messages=[{"role": "user", "content": json_prompt}],
|
||||
temperature=kwargs.get("temperature", 0.3),
|
||||
)
|
||||
# Groq requires 'json' in the prompt for json_object mode
|
||||
json_prompt = prompt
|
||||
if "json" not in prompt.lower():
|
||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": json_prompt}],
|
||||
"temperature": kwargs.get("temperature", 0.3),
|
||||
"response_format": {"type": "json_object"},
|
||||
}
|
||||
|
||||
# Support max_tokens and max_completion_tokens
|
||||
if "max_completion_tokens" in kwargs:
|
||||
create_kwargs["max_completion_tokens"] = kwargs["max_completion_tokens"]
|
||||
elif "max_tokens" in kwargs:
|
||||
create_kwargs["max_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "user"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
verbose_mode = kwargs.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [GroqProvider.generate_structured] Sending structured request to Groq API (model: {create_kwargs['model']})...", flush=True, file=sys.stdout)
|
||||
|
||||
response = self.client.chat.completions.create(**create_kwargs)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [GroqProvider.generate_structured] Structured response received from Groq.", flush=True, file=sys.stdout)
|
||||
try:
|
||||
return self._parse_json(response.choices[0].message.content)
|
||||
except Exception as e:
|
||||
@@ -469,11 +876,23 @@ class AnthropicProvider(BaseProvider):
|
||||
"Anthropic client not initialized. Set ANTHROPIC_API_KEY or pass api_key."
|
||||
)
|
||||
|
||||
response = self.client.messages.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
max_tokens=kwargs.get("max_tokens", 4096),
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
)
|
||||
# Anthropic requires max_tokens.
|
||||
# We rely on kwargs, but fallback to 8192 (safe max for newer models) if not provided.
|
||||
max_tokens = kwargs.get("max_tokens", 8192)
|
||||
|
||||
# Prepare arguments
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["temperature", "top_p", "top_k", "stop_sequences", "system", "metadata"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
response = self.client.messages.create(**create_kwargs)
|
||||
return response.content[0].text
|
||||
|
||||
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||
@@ -482,11 +901,23 @@ class AnthropicProvider(BaseProvider):
|
||||
raise ProcessingError("Anthropic client not initialized.")
|
||||
|
||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||
response = self.client.messages.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
max_tokens=kwargs.get("max_tokens", 4096),
|
||||
messages=[{"role": "user", "content": json_prompt}],
|
||||
)
|
||||
|
||||
# Anthropic requires max_tokens.
|
||||
max_tokens = kwargs.get("max_tokens", 8192)
|
||||
|
||||
# Prepare arguments
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [{"role": "user", "content": json_prompt}],
|
||||
}
|
||||
|
||||
# Pass through other common parameters
|
||||
for param in ["temperature", "top_p", "top_k", "stop_sequences", "system", "metadata"]:
|
||||
if param in kwargs:
|
||||
create_kwargs[param] = kwargs[param]
|
||||
|
||||
response = self.client.messages.create(**create_kwargs)
|
||||
try:
|
||||
return self._parse_json(response.content[0].text)
|
||||
except Exception as e:
|
||||
@@ -537,10 +968,25 @@ class OllamaProvider(BaseProvider):
|
||||
"Ollama client not initialized. Make sure Ollama is running."
|
||||
)
|
||||
|
||||
options = {"temperature": kwargs.get("temperature", 0.3)}
|
||||
|
||||
if "max_tokens" in kwargs:
|
||||
options["num_predict"] = kwargs["max_tokens"]
|
||||
|
||||
if "num_ctx" in kwargs:
|
||||
options["num_ctx"] = kwargs["num_ctx"]
|
||||
elif "context_window" in kwargs:
|
||||
options["num_ctx"] = kwargs["context_window"]
|
||||
|
||||
# Pass through other common options
|
||||
for param in ["top_p", "top_k", "repeat_penalty", "seed"]:
|
||||
if param in kwargs:
|
||||
options[param] = kwargs[param]
|
||||
|
||||
response = self.client.generate(
|
||||
model=kwargs.get("model", self.model),
|
||||
prompt=prompt,
|
||||
options={"temperature": kwargs.get("temperature", 0.3)},
|
||||
options=options,
|
||||
)
|
||||
return response.get("response", "")
|
||||
|
||||
@@ -550,10 +996,26 @@ class OllamaProvider(BaseProvider):
|
||||
raise ProcessingError("Ollama client not initialized.")
|
||||
|
||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||
|
||||
options = {"temperature": kwargs.get("temperature", 0.3)}
|
||||
|
||||
if "max_tokens" in kwargs:
|
||||
options["num_predict"] = kwargs["max_tokens"]
|
||||
|
||||
if "num_ctx" in kwargs:
|
||||
options["num_ctx"] = kwargs["num_ctx"]
|
||||
elif "context_window" in kwargs:
|
||||
options["num_ctx"] = kwargs["context_window"]
|
||||
|
||||
# Pass through other common options
|
||||
for param in ["top_p", "top_k", "repeat_penalty", "seed"]:
|
||||
if param in kwargs:
|
||||
options[param] = kwargs[param]
|
||||
|
||||
response = self.client.generate(
|
||||
model=kwargs.get("model", self.model),
|
||||
prompt=json_prompt,
|
||||
options={"temperature": kwargs.get("temperature", 0.3)},
|
||||
options=options,
|
||||
)
|
||||
try:
|
||||
return self._parse_json(response.get("response", "{}"))
|
||||
@@ -587,11 +1049,16 @@ class DeepSeekProvider(BaseProvider):
|
||||
def generate(self, prompt: str, **kwargs) -> str:
|
||||
if not self.client:
|
||||
raise ProcessingError("DeepSeek client not initialized. Set DEEPSEEK_API_KEY or pass api_key.")
|
||||
response = self.client.chat.completions.create(
|
||||
model=kwargs.get("model", self.model),
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=kwargs.get("temperature", 0.3),
|
||||
)
|
||||
|
||||
create_kwargs = {
|
||||
"model": kwargs.get("model", self.model),
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": kwargs.get("temperature", 0.3),
|
||||
}
|
||||
if "max_tokens" in kwargs:
|
||||
create_kwargs["max_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
response = self.client.chat.completions.create(**create_kwargs)
|
||||
return response.choices[0].message.content
|
||||
def generate_structured(self, prompt: str, **kwargs) -> Union[dict, list]:
|
||||
"""Generate structured output."""
|
||||
@@ -658,11 +1125,27 @@ class HuggingFaceLLMProvider(BaseProvider):
|
||||
raise ProcessingError("HuggingFace model not initialized.")
|
||||
|
||||
inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
|
||||
|
||||
# Use max_new_tokens if available, otherwise fallback to max_length with a safe default
|
||||
generate_kwargs = {
|
||||
"temperature": kwargs.get("temperature", 0.7),
|
||||
"do_sample": True,
|
||||
}
|
||||
|
||||
if "max_new_tokens" in kwargs:
|
||||
generate_kwargs["max_new_tokens"] = kwargs["max_new_tokens"]
|
||||
elif "max_tokens" in kwargs:
|
||||
generate_kwargs["max_new_tokens"] = kwargs["max_tokens"]
|
||||
|
||||
# Support legacy max_length if explicitly provided
|
||||
if "max_length" in kwargs:
|
||||
generate_kwargs["max_length"] = kwargs["max_length"]
|
||||
# Remove max_new_tokens if max_length is set to avoid conflict
|
||||
generate_kwargs.pop("max_new_tokens", None)
|
||||
|
||||
outputs = self.model.generate(
|
||||
inputs,
|
||||
max_length=kwargs.get("max_length", 100),
|
||||
temperature=kwargs.get("temperature", 0.7),
|
||||
do_sample=True,
|
||||
**generate_kwargs
|
||||
)
|
||||
generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
||||
# Remove the original prompt from the response
|
||||
@@ -786,16 +1269,33 @@ class HuggingFaceModelLoader:
|
||||
# This would need to be customized based on the model architecture
|
||||
return model(text)
|
||||
|
||||
def extract_triplets(self, model, text: str) -> List[Dict]:
|
||||
def extract_triplets(self, model, text: str, **kwargs) -> List[Dict]:
|
||||
"""Extract triplets using loaded model."""
|
||||
tokenizer = model["tokenizer"]
|
||||
model_obj = model["model"]
|
||||
device = model["device"]
|
||||
|
||||
# Use kwargs for max_length, default to 512 for input and 128 for output if not specified
|
||||
max_input_length = kwargs.get("max_input_length", 512)
|
||||
max_length = kwargs.get("max_length", 128)
|
||||
|
||||
# Allow max_new_tokens as well
|
||||
generate_kwargs = {"max_length": max_length}
|
||||
if "max_new_tokens" in kwargs:
|
||||
generate_kwargs["max_new_tokens"] = kwargs["max_new_tokens"]
|
||||
# If max_new_tokens is set, we might want to remove max_length or ensure they don't conflict
|
||||
# For Seq2Seq, max_length usually refers to the total length of the target sequence
|
||||
|
||||
# Pass other generation args
|
||||
for param in ["num_beams", "temperature", "top_p", "top_k", "do_sample"]:
|
||||
if param in kwargs:
|
||||
generate_kwargs[param] = kwargs[param]
|
||||
|
||||
inputs = tokenizer(
|
||||
text, return_tensors="pt", truncation=True, max_length=512
|
||||
text, return_tensors="pt", truncation=True, max_length=max_input_length
|
||||
).to(device)
|
||||
outputs = model_obj.generate(**inputs, max_length=128)
|
||||
|
||||
outputs = model_obj.generate(**inputs, **generate_kwargs)
|
||||
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
||||
|
||||
# Parse decoded output (format depends on model)
|
||||
@@ -803,28 +1303,88 @@ class HuggingFaceModelLoader:
|
||||
return [{"triplet": decoded}]
|
||||
|
||||
|
||||
def create_provider(name: str, **kwargs) -> BaseProvider:
|
||||
"""Create provider - checks registry for custom providers."""
|
||||
# Check registry first
|
||||
custom_provider = provider_registry.get(name)
|
||||
if custom_provider:
|
||||
return custom_provider(**kwargs)
|
||||
class ProviderPool:
|
||||
"""Pool for reusing provider instances."""
|
||||
|
||||
def __init__(self):
|
||||
self._providers: Dict[str, BaseProvider] = {}
|
||||
self.logger = get_logger("provider_pool")
|
||||
|
||||
# Built-in providers
|
||||
builtin = {
|
||||
"openai": OpenAIProvider,
|
||||
"gemini": GeminiProvider,
|
||||
"groq": GroqProvider,
|
||||
"anthropic": AnthropicProvider,
|
||||
"ollama": OllamaProvider,
|
||||
"huggingface_llm": HuggingFaceLLMProvider,
|
||||
"deepseek": DeepSeekProvider,
|
||||
}
|
||||
def get(self, name: str, **kwargs) -> BaseProvider:
|
||||
"""Get or create a provider instance."""
|
||||
# Create a cache key from name and kwargs
|
||||
# Filter out non-hashable items or volatile args if any
|
||||
# For now, we assume kwargs are configuration options that should match
|
||||
|
||||
# Helper to make dict hashable
|
||||
def make_hashable(value):
|
||||
if isinstance(value, dict):
|
||||
return tuple(sorted((k, make_hashable(v)) for k, v in value.items()))
|
||||
elif isinstance(value, list):
|
||||
return tuple(make_hashable(v) for v in value)
|
||||
return value
|
||||
|
||||
provider_class = builtin.get(name.lower())
|
||||
if not provider_class:
|
||||
raise ValueError(
|
||||
f"Unknown provider: {name}. Register custom provider or use built-in: {list(builtin.keys())}"
|
||||
)
|
||||
key_parts = [name]
|
||||
for k, v in sorted(kwargs.items()):
|
||||
# Skip some keys if they shouldn't affect pooling?
|
||||
# For now, all init args matter for the instance identity.
|
||||
key_parts.append((k, make_hashable(v)))
|
||||
|
||||
key = str(tuple(key_parts))
|
||||
|
||||
if key in self._providers:
|
||||
return self._providers[key]
|
||||
|
||||
self.logger.debug(f"Creating new provider instance for {name}")
|
||||
provider = self._create_provider(name, **kwargs)
|
||||
self._providers[key] = provider
|
||||
return provider
|
||||
|
||||
def _create_provider(self, name: str, **kwargs) -> BaseProvider:
|
||||
"""Internal creation logic."""
|
||||
# Check registry first
|
||||
custom_provider = provider_registry.get(name)
|
||||
if custom_provider:
|
||||
return custom_provider(**kwargs)
|
||||
|
||||
return provider_class(**kwargs)
|
||||
# Built-in providers
|
||||
builtin = {
|
||||
"openai": OpenAIProvider,
|
||||
"gemini": GeminiProvider,
|
||||
"groq": GroqProvider,
|
||||
"anthropic": AnthropicProvider,
|
||||
"ollama": OllamaProvider,
|
||||
"huggingface_llm": HuggingFaceLLMProvider,
|
||||
"deepseek": DeepSeekProvider,
|
||||
}
|
||||
|
||||
provider_class = builtin.get(name.lower())
|
||||
if not provider_class:
|
||||
raise ValueError(
|
||||
f"Unknown provider: {name}. Register custom provider or use built-in: {list(builtin.keys())}"
|
||||
)
|
||||
|
||||
return provider_class(**kwargs)
|
||||
|
||||
def clear(self):
|
||||
"""Clear the provider pool."""
|
||||
self._providers.clear()
|
||||
|
||||
|
||||
# Global provider pool
|
||||
_provider_pool = ProviderPool()
|
||||
|
||||
|
||||
def create_provider(name: str, use_pool: bool = True, **kwargs) -> BaseProvider:
|
||||
"""
|
||||
Create provider - checks registry for custom providers.
|
||||
|
||||
Args:
|
||||
name: Provider name
|
||||
use_pool: Whether to use the provider pool (default: True)
|
||||
**kwargs: Provider arguments
|
||||
"""
|
||||
if use_pool:
|
||||
return _provider_pool.get(name, **kwargs)
|
||||
|
||||
return _provider_pool._create_provider(name, **kwargs)
|
||||
|
||||
@@ -20,6 +20,12 @@ Algorithms Used:
|
||||
- Sequence Classification: Transformer-based relation classification models
|
||||
- Large Language Models: GPT, Claude, Gemini for relation extraction
|
||||
- Context Window Analysis: Sliding window and context extraction algorithms
|
||||
- Weighted Confidence Scoring:
|
||||
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||
* Method_Confidence: Confidence score from the extraction algorithm
|
||||
* Type_Similarity_Score: Semantic match with user-provided relation types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||
- Last Resort Fallback: Adjacency-based heuristic when all other methods fail
|
||||
|
||||
Key Features:
|
||||
- Multiple extraction methods:
|
||||
@@ -30,6 +36,7 @@ Key Features:
|
||||
* HuggingFace: Custom HuggingFace relation models
|
||||
* LLM-based: LLM-powered relation extraction
|
||||
- Fallback chain support: Try methods in order until one succeeds
|
||||
- Robust Fallbacks: Prevents empty results via Primary -> Pattern -> Last Resort chain
|
||||
- Multiple relation types (founded_by, located_in, works_for, born_in, etc.)
|
||||
- Relation classification and grouping
|
||||
- Relation validation and consistency checking
|
||||
@@ -194,9 +201,12 @@ class RelationExtractor:
|
||||
)
|
||||
|
||||
try:
|
||||
results = []
|
||||
# Ensure lists are same length
|
||||
min_len = min(len(text), len(entities))
|
||||
results = [None] * min_len
|
||||
total_relations_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
||||
if min_len <= 10:
|
||||
update_interval = 1 # Update every item for small datasets
|
||||
@@ -204,52 +214,106 @@ class RelationExtractor:
|
||||
update_interval = max(1, min(10, min_len // 100))
|
||||
|
||||
# Initial progress update - ALWAYS show this
|
||||
remaining = min_len
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=min_len,
|
||||
message=f"Starting batch extraction... 0/{min_len} (remaining: {remaining})"
|
||||
message=f"Starting batch extraction... 0/{min_len}"
|
||||
)
|
||||
|
||||
for i in range(min_len):
|
||||
doc_item = text[i]
|
||||
ent_item = entities[i]
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
def process_item(i, doc_item, ent_item):
|
||||
try:
|
||||
doc_text = ""
|
||||
if isinstance(doc_item, dict) and "content" in doc_item:
|
||||
doc_text = doc_item["content"]
|
||||
elif isinstance(doc_item, str):
|
||||
doc_text = doc_item
|
||||
else:
|
||||
doc_text = str(doc_item)
|
||||
|
||||
# Ensure ent_item is a list of entities
|
||||
if not isinstance(ent_item, list):
|
||||
ent_item = [] # Should not happen if entities is List[List[Entity]]
|
||||
|
||||
current_relations = self.extract_relations(doc_text, ent_item, **kwargs)
|
||||
|
||||
# Add provenance metadata
|
||||
for rel in current_relations:
|
||||
if rel.metadata is None:
|
||||
rel.metadata = {}
|
||||
rel.metadata["batch_index"] = i
|
||||
if isinstance(doc_item, dict) and "id" in doc_item:
|
||||
rel.metadata["document_id"] = doc_item["id"]
|
||||
|
||||
return i, current_relations
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to process item {i}: {e}")
|
||||
return i, []
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
doc_text = ""
|
||||
if isinstance(doc_item, dict) and "content" in doc_item:
|
||||
doc_text = doc_item["content"]
|
||||
elif isinstance(doc_item, str):
|
||||
doc_text = doc_item
|
||||
else:
|
||||
doc_text = str(doc_item)
|
||||
|
||||
# Ensure ent_item is a list of entities
|
||||
if not isinstance(ent_item, list):
|
||||
ent_item = [] # Should not happen if entities is List[List[Entity]]
|
||||
|
||||
results.append(self.extract_relations(doc_text, ent_item, **kwargs))
|
||||
|
||||
remaining = min_len - (i + 1)
|
||||
# Update progress: always update for small datasets, or at intervals for large ones
|
||||
should_update = (
|
||||
(i + 1) % update_interval == 0 or
|
||||
(i + 1) == min_len or
|
||||
i == 0 or
|
||||
min_len <= 10 # Always update for small datasets
|
||||
)
|
||||
if should_update:
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=i + 1,
|
||||
total=min_len,
|
||||
message=f"Processing documents... {i + 1}/{min_len} (remaining: {remaining})"
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit tasks
|
||||
future_to_idx = {
|
||||
executor.submit(process_item, i, text[i], entities[i]): i
|
||||
for i in range(min_len)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
i, relations = future.result()
|
||||
results[i] = relations
|
||||
total_relations_count += len(relations)
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0 or
|
||||
processed_count == min_len or
|
||||
processed_count == 1 or
|
||||
min_len <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = min_len - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=min_len,
|
||||
message=f"Processing documents... {processed_count}/{min_len} (remaining: {remaining}) - Extracted {total_relations_count} relations so far"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
for i in range(min_len):
|
||||
_, relations = process_item(i, text[i], entities[i])
|
||||
results[i] = relations
|
||||
total_relations_count += len(relations)
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0 or
|
||||
processed_count == min_len or
|
||||
processed_count == 1 or
|
||||
min_len <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = min_len - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=min_len,
|
||||
message=f"Processing documents... {processed_count}/{min_len} (remaining: {remaining}) - Extracted {total_relations_count} relations so far"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Extracted relations from {len(results)} documents",
|
||||
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_relations_count} relations.",
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
@@ -265,14 +329,19 @@ class RelationExtractor:
|
||||
return []
|
||||
|
||||
def extract_relations(
|
||||
self, text: str, entities: List[Entity], **options
|
||||
) -> List[Relation]:
|
||||
self,
|
||||
text: Union[str, List[Dict[str, Any]], List[str]],
|
||||
entities: Union[List[Entity], List[List[Entity]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> Union[List[Relation], List[List[Relation]]]:
|
||||
"""
|
||||
Extract relations between entities.
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
entities: List of extracted entities
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options:
|
||||
- method: Override method (if not set in __init__)
|
||||
- min_confidence: Minimum confidence threshold
|
||||
@@ -281,7 +350,18 @@ class RelationExtractor:
|
||||
Returns:
|
||||
list: List of extracted relations
|
||||
"""
|
||||
from .methods import get_relation_method
|
||||
if isinstance(text, list):
|
||||
if entities is None:
|
||||
entities_batch = [[] for _ in range(len(text))]
|
||||
elif isinstance(entities, list) and (not entities):
|
||||
entities_batch = [[] for _ in range(len(text))]
|
||||
elif isinstance(entities, list) and all(isinstance(e, Entity) for e in entities):
|
||||
entities_batch = [entities for _ in range(len(text))]
|
||||
else:
|
||||
entities_batch = entities
|
||||
return self.extract(text, entities_batch, pipeline_id=pipeline_id, **options)
|
||||
|
||||
from .methods import get_relation_method, match_entity
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
@@ -322,6 +402,11 @@ class RelationExtractor:
|
||||
|
||||
# Prepare method-specific options
|
||||
method_options = all_options.copy()
|
||||
|
||||
# Pass relation_types to all methods so they can use them (e.g. for similarity matching)
|
||||
if relation_types:
|
||||
method_options["relation_types"] = relation_types
|
||||
|
||||
if method_name == "huggingface":
|
||||
method_options["model"] = all_options.get(
|
||||
"huggingface_model", all_options.get("model")
|
||||
@@ -334,20 +419,16 @@ class RelationExtractor:
|
||||
method_options["model"] = all_options.get(
|
||||
"llm_model", all_options.get("model")
|
||||
)
|
||||
# Pass api_key if provided (needed for all providers)
|
||||
if "api_key" in all_options:
|
||||
method_options["api_key"] = all_options["api_key"]
|
||||
elif "api_key" not in method_options:
|
||||
# Try to get from environment as fallback
|
||||
# Ensure api_key is populated: check explicitly provided or fallback to env
|
||||
current_key = method_options.get("api_key")
|
||||
if not current_key:
|
||||
# Not found or empty/None, try environment
|
||||
import os
|
||||
provider = method_options.get("provider", "openai")
|
||||
env_key = f"{provider.upper()}_API_KEY"
|
||||
provider_name = method_options.get("provider", "openai")
|
||||
env_key = f"{provider_name.upper()}_API_KEY"
|
||||
api_key = os.getenv(env_key)
|
||||
if api_key:
|
||||
method_options["api_key"] = api_key
|
||||
# Pass relation_types to LLM method so it can use them in the prompt
|
||||
if relation_types:
|
||||
method_options["relation_types"] = relation_types
|
||||
elif method_name == "dependency":
|
||||
method_options["model"] = all_options.get(
|
||||
"model", "en_core_web_sm"
|
||||
@@ -358,6 +439,12 @@ class RelationExtractor:
|
||||
if verbose_mode and method_name == "llm":
|
||||
import sys
|
||||
print(f" [RelationExtractor] Processing with {method_name}...", flush=True, file=sys.stdout)
|
||||
print(f" [RelationExtractor Debug] method_options keys: {list(method_options.keys())}", flush=True, file=sys.stdout)
|
||||
if "api_key" in method_options:
|
||||
masked = method_options["api_key"][:4] + "..." if method_options["api_key"] else "None"
|
||||
print(f" [RelationExtractor Debug] api_key present: {masked}", flush=True, file=sys.stdout)
|
||||
else:
|
||||
print(f" [RelationExtractor Debug] api_key NOT present", flush=True, file=sys.stdout)
|
||||
|
||||
relations = method_func(text, entities, **method_options)
|
||||
|
||||
@@ -366,6 +453,20 @@ class RelationExtractor:
|
||||
import sys
|
||||
print(f" [RelationExtractor] Extracted {len(relations)} relations", flush=True, file=sys.stdout)
|
||||
|
||||
# Apply weighted scoring if relation_types are provided
|
||||
if relation_types:
|
||||
try:
|
||||
from .methods import calculate_weighted_confidence
|
||||
for r in relations:
|
||||
r.confidence = calculate_weighted_confidence(
|
||||
item_type=r.predicate,
|
||||
original_confidence=r.confidence,
|
||||
valid_types=relation_types,
|
||||
item_text=r.predicate # For relations, the predicate IS the text usually
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Filter by confidence
|
||||
filtered = [r for r in relations if r.confidence >= min_confidence]
|
||||
|
||||
@@ -386,13 +487,23 @@ class RelationExtractor:
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Method {method_name} failed: {e}")
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [RelationExtractor] ERROR: Method {method_name} failed: {e}", flush=True, file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
continue
|
||||
|
||||
# Use first successful method or combine
|
||||
if all_relations:
|
||||
relations = all_relations[0][1] # Use first successful method
|
||||
else:
|
||||
relations = []
|
||||
# Fallback to pattern-based extraction if all models fail
|
||||
relations = self._extract_with_patterns(text, entities)
|
||||
|
||||
# Last resort: if patterns also fail but we have entities, force some relations
|
||||
if not relations and entities and len(entities) >= 2:
|
||||
relations = self._extract_last_resort_relations(text, entities)
|
||||
|
||||
# Validate if enabled
|
||||
if validate:
|
||||
@@ -411,15 +522,45 @@ class RelationExtractor:
|
||||
)
|
||||
raise
|
||||
|
||||
def _extract_last_resort_relations(self, text: str, entities: List[Entity]) -> List[Relation]:
|
||||
"""Last resort relation extraction based on simple adjacency."""
|
||||
relations = []
|
||||
# Connect adjacent entities
|
||||
for i in range(len(entities) - 1):
|
||||
e1 = entities[i]
|
||||
e2 = entities[i+1]
|
||||
|
||||
# Create a weak relation
|
||||
start_idx = min(e1.end_char, e2.start_char)
|
||||
end_idx = max(e1.end_char, e2.start_char)
|
||||
# Ensure context isn't too large or invalid
|
||||
if start_idx < 0: start_idx = 0
|
||||
if end_idx > len(text): end_idx = len(text)
|
||||
|
||||
# Expand context a bit
|
||||
ctx_start = max(0, start_idx - 20)
|
||||
ctx_end = min(len(text), end_idx + 20)
|
||||
|
||||
context = text[ctx_start:ctx_end]
|
||||
|
||||
rel = Relation(
|
||||
subject=e1,
|
||||
predicate="related_to",
|
||||
object=e2,
|
||||
confidence=0.3,
|
||||
context=context,
|
||||
metadata={"extraction_method": "last_resort_adjacency"}
|
||||
)
|
||||
relations.append(rel)
|
||||
return relations
|
||||
|
||||
def _extract_with_patterns(
|
||||
self, text: str, entities: List[Entity]
|
||||
) -> List[Relation]:
|
||||
"""Extract relations using pattern matching."""
|
||||
from .methods import match_entity
|
||||
relations = []
|
||||
|
||||
# Create entity lookup by text
|
||||
entity_map = {e.text.lower(): e for e in entities}
|
||||
|
||||
# Check each relation pattern
|
||||
for relation_type, patterns in self.relation_patterns.items():
|
||||
for pattern in patterns:
|
||||
@@ -427,8 +568,8 @@ class RelationExtractor:
|
||||
subject_text = match.group("subject").strip()
|
||||
object_text = match.group("object").strip()
|
||||
|
||||
subject_entity = entity_map.get(subject_text.lower())
|
||||
object_entity = entity_map.get(object_text.lower())
|
||||
subject_entity = match_entity(subject_text, entities)
|
||||
object_entity = match_entity(object_text, entities)
|
||||
|
||||
if subject_entity and object_entity:
|
||||
# Get context around the match
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict
|
||||
|
||||
class EntityOut(BaseModel):
|
||||
"""Canonical schema for entity extraction output."""
|
||||
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||
|
||||
text: str = Field(..., description="The text content of the entity")
|
||||
label: str = Field(..., description="The type or label of the entity (e.g., PERSON, ORG)")
|
||||
start: int = Field(0, description="Start character index", alias="start_char")
|
||||
end: int = Field(0, description="End character index", alias="end_char")
|
||||
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||
|
||||
@field_validator("text", mode="before")
|
||||
@classmethod
|
||||
def clean_text(cls, v):
|
||||
if isinstance(v, str):
|
||||
return v.strip()
|
||||
return str(v)
|
||||
|
||||
@field_validator("confidence", mode="before")
|
||||
@classmethod
|
||||
def normalize_confidence(cls, v):
|
||||
if isinstance(v, str):
|
||||
try:
|
||||
v = float(v)
|
||||
except ValueError:
|
||||
return 0.9
|
||||
if isinstance(v, (int, float)):
|
||||
return max(0.0, min(1.0, float(v)))
|
||||
return 0.9
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def handle_aliases(cls, data):
|
||||
if isinstance(data, dict):
|
||||
# Handle 'type' as alias for 'label'
|
||||
if "label" not in data and "type" in data:
|
||||
data["label"] = data["type"]
|
||||
# Handle 'value' or 'span' as alias for 'text'
|
||||
if "text" not in data:
|
||||
if "value" in data:
|
||||
data["text"] = data["value"]
|
||||
elif "span" in data:
|
||||
data["text"] = data["span"]
|
||||
return data
|
||||
|
||||
class RelationOut(BaseModel):
|
||||
"""Canonical schema for relation extraction output."""
|
||||
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||
|
||||
subject: str = Field(..., description="Source entity text")
|
||||
object: str = Field(..., description="Target entity text")
|
||||
predicate: str = Field(..., description="Relation type or predicate")
|
||||
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def handle_aliases(cls, data):
|
||||
if isinstance(data, dict):
|
||||
if "subject" not in data and "source" in data:
|
||||
data["subject"] = data["source"]
|
||||
if "object" not in data and "target" in data:
|
||||
data["object"] = data["target"]
|
||||
if "predicate" not in data and "label" in data:
|
||||
data["predicate"] = data["label"]
|
||||
return data
|
||||
|
||||
@property
|
||||
def source(self):
|
||||
return self.subject
|
||||
|
||||
@property
|
||||
def target(self):
|
||||
return self.object
|
||||
|
||||
@property
|
||||
def label(self):
|
||||
return self.predicate
|
||||
|
||||
@field_validator("confidence", mode="before")
|
||||
@classmethod
|
||||
def normalize_confidence(cls, v):
|
||||
if isinstance(v, str):
|
||||
try:
|
||||
v = float(v)
|
||||
except ValueError:
|
||||
return 0.9
|
||||
if isinstance(v, (int, float)):
|
||||
return max(0.0, min(1.0, float(v)))
|
||||
return 0.9
|
||||
|
||||
class TripletOut(BaseModel):
|
||||
"""Canonical schema for triplet extraction output."""
|
||||
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||
|
||||
subject: str = Field(..., description="Subject of the triplet")
|
||||
predicate: str = Field(..., description="Predicate or relation")
|
||||
object: str = Field(..., description="Object of the triplet")
|
||||
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||
|
||||
@field_validator("confidence", mode="before")
|
||||
@classmethod
|
||||
def normalize_confidence(cls, v):
|
||||
if isinstance(v, str):
|
||||
try:
|
||||
v = float(v)
|
||||
except ValueError:
|
||||
return 0.9
|
||||
if isinstance(v, (int, float)):
|
||||
return max(0.0, min(1.0, float(v)))
|
||||
return 0.9
|
||||
|
||||
class EntitiesResponse(BaseModel):
|
||||
"""Wrapper for list of entities."""
|
||||
entities: List[EntityOut] = Field(default_factory=list)
|
||||
|
||||
class RelationsResponse(BaseModel):
|
||||
"""Wrapper for list of relations."""
|
||||
relations: List[RelationOut] = Field(default_factory=list)
|
||||
|
||||
class TripletsResponse(BaseModel):
|
||||
"""Wrapper for list of triplets."""
|
||||
triplets: List[TripletOut] = Field(default_factory=list)
|
||||
@@ -71,6 +71,7 @@ class SemanticRole:
|
||||
start_char: int
|
||||
end_char: int
|
||||
confidence: float = 1.0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -81,6 +82,7 @@ class SemanticCluster:
|
||||
cluster_id: int
|
||||
centroid: Optional[str] = None
|
||||
similarity_score: float = 0.0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class SemanticAnalyzer:
|
||||
@@ -118,6 +120,147 @@ class SemanticAnalyzer:
|
||||
self.role_labeler = RoleLabeler(**self.config.get("role", {}))
|
||||
self.semantic_clusterer = SemanticClusterer(**self.config.get("clustering", {}))
|
||||
|
||||
def analyze(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
pipeline_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
|
||||
"""
|
||||
Perform semantic analysis on text or list of documents.
|
||||
Handles batch processing with progress tracking.
|
||||
|
||||
Args:
|
||||
text: Input text or list of documents
|
||||
pipeline_id: Optional pipeline ID for progress tracking
|
||||
**kwargs: Analysis options
|
||||
|
||||
Returns:
|
||||
Union[Dict[str, Any], List[Dict[str, Any]]]: Analysis results
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
# Handle batch analysis with progress tracking
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="SemanticAnalyzer",
|
||||
message=f"Batch analyzing {len(text)} documents",
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
try:
|
||||
results = [None] * len(text)
|
||||
total_items = len(text)
|
||||
processed_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch analysis... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
)
|
||||
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
analysis = self.analyze_semantics(doc_text, **kwargs)
|
||||
|
||||
analysis["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
analysis["document_id"] = item["id"]
|
||||
|
||||
if "semantic_roles" in analysis:
|
||||
for role in analysis["semantic_roles"]:
|
||||
if "metadata" not in role:
|
||||
role["metadata"] = {}
|
||||
|
||||
role["metadata"]["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
role["metadata"]["document_id"] = item["id"]
|
||||
|
||||
return idx, analysis
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to analyze item {idx}: {e}")
|
||||
return idx, {"error": str(e), "batch_index": idx}
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_idx = {
|
||||
executor.submit(process_item, idx, item): idx
|
||||
for idx, item in enumerate(text)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, analysis = future.result()
|
||||
results[idx] = analysis
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0
|
||||
or processed_count == total_items
|
||||
or processed_count == 1
|
||||
or total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
else:
|
||||
for idx, item in enumerate(text):
|
||||
_, analysis = process_item(idx, item)
|
||||
results[idx] = analysis
|
||||
processed_count += 1
|
||||
|
||||
should_update = (
|
||||
processed_count % update_interval == 0
|
||||
or processed_count == total_items
|
||||
or processed_count == 1
|
||||
or total_items <= 10
|
||||
)
|
||||
if should_update:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch analysis completed. Processed {len(results)} documents.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
else:
|
||||
# Single item
|
||||
return self.analyze_semantics(text, **kwargs)
|
||||
|
||||
def analyze_semantics(self, text: str, **options) -> Dict[str, Any]:
|
||||
"""
|
||||
Perform comprehensive semantic analysis.
|
||||
@@ -200,13 +343,13 @@ class SemanticAnalyzer:
|
||||
return self.label_semantic_roles(text, **options)
|
||||
|
||||
def cluster_semantically(
|
||||
self, texts: List[str], **options
|
||||
self, texts: Union[List[str], List[Dict[str, Any]]], **options
|
||||
) -> List[SemanticCluster]:
|
||||
"""
|
||||
Perform semantic clustering of texts.
|
||||
|
||||
Args:
|
||||
texts: List of texts to cluster
|
||||
texts: List of texts or documents to cluster
|
||||
**options: Clustering options
|
||||
|
||||
Returns:
|
||||
@@ -372,12 +515,14 @@ class SemanticClusterer:
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
def cluster(self, texts: List[str], **options) -> List[SemanticCluster]:
|
||||
def cluster(
|
||||
self, texts: Union[List[str], List[Dict[str, Any]]], **options
|
||||
) -> List[SemanticCluster]:
|
||||
"""
|
||||
Perform semantic clustering of texts.
|
||||
|
||||
Args:
|
||||
texts: List of texts to cluster
|
||||
texts: List of texts or documents (dict with 'content' and 'id') to cluster
|
||||
**options: Clustering options:
|
||||
- num_clusters: Number of clusters (default: auto)
|
||||
- similarity_threshold: Minimum similarity for clustering
|
||||
@@ -388,11 +533,27 @@ class SemanticClusterer:
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
# Extract content and IDs if input is list of dicts
|
||||
processed_texts = []
|
||||
doc_ids = []
|
||||
|
||||
for item in texts:
|
||||
if isinstance(item, dict):
|
||||
content = item.get("content", str(item))
|
||||
processed_texts.append(content)
|
||||
if "id" in item:
|
||||
doc_ids.append(item["id"])
|
||||
else:
|
||||
doc_ids.append(None)
|
||||
else:
|
||||
processed_texts.append(str(item))
|
||||
doc_ids.append(None)
|
||||
|
||||
# Track clustering
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="SemanticClusterer",
|
||||
message=f"Clustering {len(texts)} texts",
|
||||
message=f"Clustering {len(processed_texts)} texts",
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -402,7 +563,7 @@ class SemanticClusterer:
|
||||
clusters = []
|
||||
assigned = set()
|
||||
|
||||
total_texts = len(texts)
|
||||
total_texts = len(processed_texts)
|
||||
if total_texts <= 10:
|
||||
update_interval = 1 # Update every item for small datasets
|
||||
else:
|
||||
@@ -418,22 +579,28 @@ class SemanticClusterer:
|
||||
)
|
||||
|
||||
cluster_id = 0
|
||||
for i, text1 in enumerate(texts):
|
||||
for i, text1 in enumerate(processed_texts):
|
||||
if i in assigned:
|
||||
continue
|
||||
|
||||
cluster_texts = [text1]
|
||||
cluster_doc_ids = []
|
||||
if doc_ids[i] is not None:
|
||||
cluster_doc_ids.append(doc_ids[i])
|
||||
|
||||
assigned.add(i)
|
||||
|
||||
# Find similar texts
|
||||
remaining_texts = len(texts) - (i + 1)
|
||||
for j, text2 in enumerate(texts[i + 1 :], start=i + 1):
|
||||
remaining_texts = len(processed_texts) - (i + 1)
|
||||
for j, text2 in enumerate(processed_texts[i + 1 :], start=i + 1):
|
||||
if j in assigned:
|
||||
continue
|
||||
|
||||
similarity = similarity_analyzer.calculate_similarity(text1, text2)
|
||||
if similarity >= similarity_threshold:
|
||||
cluster_texts.append(text2)
|
||||
if doc_ids[j] is not None:
|
||||
cluster_doc_ids.append(doc_ids[j])
|
||||
assigned.add(j)
|
||||
|
||||
# Create cluster
|
||||
@@ -443,6 +610,11 @@ class SemanticClusterer:
|
||||
centroid=cluster_texts[0], # Use first as centroid
|
||||
similarity_score=similarity_threshold,
|
||||
)
|
||||
|
||||
# Add provenance metadata
|
||||
if cluster_doc_ids:
|
||||
cluster.metadata["document_ids"] = cluster_doc_ids
|
||||
|
||||
clusters.append(cluster)
|
||||
cluster_id += 1
|
||||
|
||||
|
||||
@@ -36,6 +36,46 @@ print(f"Relations: {relations}")
|
||||
print(f"Extracted {len(entities)} entities and {len(relations)} relations")
|
||||
```
|
||||
|
||||
## Batch Processing & Provenance
|
||||
|
||||
All extractors support batch processing for high-throughput extraction. You can pass a list of strings or a list of dictionaries (with `content` and `id` keys).
|
||||
|
||||
**Features:**
|
||||
- **Parallel Processing**: Multi-threaded extraction for high throughput (control via `max_workers`).
|
||||
- **Progress Tracking**: Automatically shows a progress bar for large batches.
|
||||
- **Provenance Metadata**: Each extracted item includes `batch_index` and `document_id` in its `metadata`.
|
||||
|
||||
```python
|
||||
from semantica.semantic_extract import NERExtractor
|
||||
|
||||
documents = [
|
||||
{"id": "doc_1", "content": "Apple Inc. was founded by Steve Jobs."},
|
||||
{"id": "doc_2", "content": "Microsoft Corporation was founded by Bill Gates."}
|
||||
]
|
||||
|
||||
# Initialize with parallel processing enabled
|
||||
extractor = NERExtractor(max_workers=4)
|
||||
batch_results = extractor.extract(documents)
|
||||
|
||||
# OR override during extraction call
|
||||
# batch_results = extractor.extract(documents, max_workers=8)
|
||||
|
||||
for i, doc_entities in enumerate(batch_results):
|
||||
print(f"Document {i} entities:")
|
||||
for entity in doc_entities:
|
||||
print(f" - {entity.text} ({entity.label})")
|
||||
print(f" Provenance: Batch Index {entity.metadata['batch_index']}, Doc ID {entity.metadata.get('document_id')}")
|
||||
```
|
||||
|
||||
## Robust Extraction Fallbacks
|
||||
|
||||
The framework implements robust fallback chains to prevent empty results when primary methods fail (e.g., due to model unavailability or obscure text).
|
||||
|
||||
- **NER**: `ML/LLM` -> `Pattern` -> `Last Resort` (Capitalized Words)
|
||||
- **Relation**: `Primary` -> `Pattern` -> `Last Resort` (Adjacency)
|
||||
- **Triplet**: `Primary` -> `Relation-to-Triplet` -> `Pattern`
|
||||
|
||||
This ensures that you almost always get *some* structured data, even if it requires falling back to simpler heuristics.
|
||||
|
||||
## Entity Extraction
|
||||
|
||||
@@ -85,9 +125,22 @@ entities = extractor.extract(
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
silent_fail=False, # Raise ProcessingError on failure (default)
|
||||
max_text_length=4000 # Auto-chunking for long text
|
||||
max_text_length=4000, # Auto-chunking for long text (default: 64k for major providers)
|
||||
max_tokens=4096, # Explicitly control generation output length
|
||||
temperature=0.0
|
||||
)
|
||||
print(f"LLM method: {len(entities)} entities")
|
||||
|
||||
# Groq extraction with long context support
|
||||
# Groq defaults to 64k chunking limit for models like llama-3.3-70b
|
||||
groq_extractor = NERExtractor(method="llm")
|
||||
groq_entities = groq_extractor.extract(
|
||||
text,
|
||||
provider="groq",
|
||||
model="llama-3.3-70b-versatile",
|
||||
max_tokens=8000 # Passed directly to Groq API
|
||||
)
|
||||
print(f"Groq method: {len(groq_entities)} entities")
|
||||
```
|
||||
|
||||
### Using NERExtractor Directly
|
||||
@@ -186,6 +239,8 @@ relations = extractor.extract(
|
||||
text,
|
||||
entities=entities,
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_tokens=2048, # Increased output limit for many relations
|
||||
silent_fail=True # Return empty list if extraction fails
|
||||
)
|
||||
```
|
||||
@@ -249,7 +304,8 @@ triplets = extractor.extract_triplets(
|
||||
text,
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_text_length=2000 # Force chunking for long text
|
||||
max_text_length=64000, # Large default chunk size supported
|
||||
max_tokens=4096 # Ensure enough tokens for all triplets
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
@@ -149,13 +149,191 @@ class SemanticNetworkExtractor:
|
||||
self.config["ner_method"] = method
|
||||
self.config["relation_method"] = method
|
||||
|
||||
self._ner_extractor = None
|
||||
self._relation_extractor = None
|
||||
|
||||
def extract(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Union[SemanticNetwork, List[SemanticNetwork]]:
|
||||
"""
|
||||
Extract semantic network from text or list of documents.
|
||||
Handles batch processing with progress tracking.
|
||||
|
||||
Args:
|
||||
text: Input text or list of documents
|
||||
entities: Optional pre-extracted entities (single list or list of lists)
|
||||
relations: Optional pre-extracted relations (single list or list of lists)
|
||||
pipeline_id: Optional pipeline ID for progress tracking
|
||||
**kwargs: Extraction options
|
||||
|
||||
Returns:
|
||||
Union[SemanticNetwork, List[SemanticNetwork]]: Extracted semantic network(s)
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
# Handle batch extraction with progress tracking
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="SemanticNetworkExtractor",
|
||||
message=f"Batch extracting semantic networks from {len(text)} documents",
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
try:
|
||||
results = [None] * len(text)
|
||||
total_items = len(text)
|
||||
processed_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch extraction... 0/{total_items} (remaining: {total_items})"
|
||||
)
|
||||
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
)
|
||||
|
||||
def process_item(idx, item, doc_entities, doc_relations):
|
||||
try:
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
|
||||
# Extract
|
||||
network = self.extract_network(
|
||||
doc_text,
|
||||
entities=doc_entities,
|
||||
relations=doc_relations,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# Add provenance metadata to nodes and edges
|
||||
batch_meta = {"batch_index": idx}
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
batch_meta["document_id"] = item["id"]
|
||||
|
||||
# Update network metadata
|
||||
network.metadata.update(batch_meta)
|
||||
|
||||
# Update nodes metadata
|
||||
for node in network.nodes:
|
||||
node.metadata.update(batch_meta)
|
||||
|
||||
# Update edges metadata
|
||||
for edge in network.edges:
|
||||
edge.metadata.update(batch_meta)
|
||||
|
||||
return idx, network
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to process item {idx}: {e}")
|
||||
verbose_mode = kwargs.get("verbose", False) or self.config.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [SemanticNetworkExtractor] ERROR: Batch item {idx} failed: {e}", flush=True, file=sys.stderr)
|
||||
return idx, None
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit tasks
|
||||
future_to_idx = {}
|
||||
for idx, item in enumerate(text):
|
||||
doc_entities = None
|
||||
if entities and isinstance(entities, list) and idx < len(entities):
|
||||
doc_entities = entities[idx]
|
||||
|
||||
doc_relations = None
|
||||
if relations and isinstance(relations, list) and idx < len(relations):
|
||||
doc_relations = relations[idx]
|
||||
|
||||
future = executor.submit(process_item, idx, item, doc_entities, doc_relations)
|
||||
future_to_idx[future] = idx
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, network = future.result()
|
||||
if network:
|
||||
results[idx] = network
|
||||
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
if (processed_count) % update_interval == 0 or (processed_count) == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
for idx, item in enumerate(text):
|
||||
doc_entities = None
|
||||
if entities and isinstance(entities, list) and idx < len(entities):
|
||||
doc_entities = entities[idx]
|
||||
|
||||
doc_relations = None
|
||||
if relations and isinstance(relations, list) and idx < len(relations):
|
||||
doc_relations = relations[idx]
|
||||
|
||||
_, network = process_item(idx, item, doc_entities, doc_relations)
|
||||
if network:
|
||||
results[idx] = network
|
||||
|
||||
processed_count += 1
|
||||
|
||||
# Update progress
|
||||
if (processed_count) % update_interval == 0 or (processed_count) == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining})"
|
||||
)
|
||||
|
||||
# Filter out None results if any failed
|
||||
results = [r for r in results if r is not None]
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch extraction completed. Processed {len(results)} documents.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
else:
|
||||
# Single item
|
||||
return self.extract_network(text, entities=entities, relations=relations, **kwargs)
|
||||
|
||||
def extract_network(
|
||||
self,
|
||||
text: str,
|
||||
entities: Optional[List[Entity]] = None,
|
||||
relations: Optional[List[Relation]] = None,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> SemanticNetwork:
|
||||
) -> Union[SemanticNetwork, List[SemanticNetwork]]:
|
||||
"""
|
||||
Extract semantic network from text.
|
||||
|
||||
@@ -168,6 +346,23 @@ class SemanticNetworkExtractor:
|
||||
Returns:
|
||||
SemanticNetwork: Extracted semantic network
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
entities_batch = entities
|
||||
if entities is not None and isinstance(entities, list) and (not entities or all(isinstance(e, Entity) for e in entities)):
|
||||
entities_batch = [entities for _ in range(len(text))] if entities else [[] for _ in range(len(text))]
|
||||
|
||||
relations_batch = relations
|
||||
if relations is not None and isinstance(relations, list) and (not relations or all(isinstance(r, Relation) for r in relations)):
|
||||
relations_batch = [relations for _ in range(len(text))] if relations else [[] for _ in range(len(text))]
|
||||
|
||||
return self.extract(
|
||||
text,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
pipeline_id=pipeline_id,
|
||||
**options,
|
||||
)
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="SemanticNetworkExtractor",
|
||||
@@ -187,15 +382,16 @@ class SemanticNetworkExtractor:
|
||||
# Pass method if specified
|
||||
if "ner_method" in self.config:
|
||||
ner_config["method"] = self.config["ner_method"]
|
||||
ner = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
entities = ner.extract_entities(text, **options)
|
||||
if self._ner_extractor is None:
|
||||
self._ner_extractor = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
entities = self._ner_extractor.extract_entities(text, **options)
|
||||
|
||||
# Extract relations if not provided
|
||||
if relations is None:
|
||||
@@ -206,15 +402,16 @@ class SemanticNetworkExtractor:
|
||||
# Pass method if specified
|
||||
if "relation_method" in self.config:
|
||||
rel_config["method"] = self.config["relation_method"]
|
||||
rel_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
relations = rel_extractor.extract_relations(text, entities, **options)
|
||||
if self._relation_extractor is None:
|
||||
self._relation_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation"]
|
||||
},
|
||||
)
|
||||
relations = self._relation_extractor.extract_relations(text, entities, **options)
|
||||
|
||||
# Build network
|
||||
total_steps = 2 # Create nodes, create edges
|
||||
@@ -241,6 +438,12 @@ class SemanticNetworkExtractor:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
verbose_mode = options.get("verbose", False) or self.config.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [SemanticNetworkExtractor] ERROR: Extraction failed: {e}", flush=True, file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
raise
|
||||
|
||||
def _build_network(
|
||||
|
||||
@@ -18,6 +18,12 @@ Algorithms Used:
|
||||
- Large Language Models: GPT, Claude, Gemini for structured triplet extraction
|
||||
- RDF Serialization: Graph serialization algorithms (Turtle, N-Triples, JSON-LD)
|
||||
- URI Normalization: String normalization and URI formatting algorithms
|
||||
- Weighted Confidence Scoring:
|
||||
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||
* Method_Confidence: Confidence score from the extraction algorithm
|
||||
* Type_Similarity_Score: Semantic match with user-provided triplet types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||
- Last Resort Fallback: Relation-to-Triplet conversion when all other methods fail
|
||||
|
||||
Key Features:
|
||||
- Multiple extraction methods:
|
||||
@@ -26,6 +32,7 @@ Key Features:
|
||||
* HuggingFace: Custom HuggingFace triplet models
|
||||
* LLM-based: LLM-powered triplet extraction
|
||||
- Fallback chain support: Try methods in order until one succeeds
|
||||
- Robust Fallbacks: Prevents empty results via Primary -> Relation-to-Triplet -> Pattern chain
|
||||
- RDF triplet generation from entities and relations
|
||||
- Subject-predicate-object extraction
|
||||
- Triplet validation and quality checking
|
||||
@@ -99,6 +106,7 @@ class TripletExtractor:
|
||||
def __init__(
|
||||
self,
|
||||
method: Union[str, List[str]] = "pattern",
|
||||
triplet_types: Optional[List[str]] = None,
|
||||
include_temporal: bool = False,
|
||||
include_provenance: bool = False,
|
||||
config=None,
|
||||
@@ -114,6 +122,7 @@ class TripletExtractor:
|
||||
- "huggingface": HuggingFace model
|
||||
- "llm": LLM-based extraction
|
||||
- List of methods for fallback chain
|
||||
triplet_types: Specific triplet types/predicates to extract (e.g., ["foundedBy", "locatedIn"])
|
||||
include_temporal: Whether to include temporal information in triplets
|
||||
include_provenance: Whether to track source sentences for provenance
|
||||
config: Legacy config dict (deprecated, use kwargs)
|
||||
@@ -134,7 +143,15 @@ class TripletExtractor:
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
if method is not None:
|
||||
self.config["ner_method"] = method
|
||||
self.config["relation_method"] = method
|
||||
|
||||
self._ner_extractor = None
|
||||
self._relation_extractor = None
|
||||
|
||||
# Store parameters
|
||||
self.triplet_types = triplet_types
|
||||
self.include_temporal = include_temporal
|
||||
self.include_provenance = include_provenance
|
||||
|
||||
@@ -149,13 +166,164 @@ class TripletExtractor:
|
||||
|
||||
self.supported_formats = ["turtle", "ntriples", "jsonld", "xml"]
|
||||
|
||||
def extract(
|
||||
self,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Union[List[Triplet], List[List[Triplet]]]:
|
||||
"""
|
||||
Extract triplets from text or list of documents.
|
||||
Handles batch processing with progress tracking.
|
||||
|
||||
Args:
|
||||
text: Input text or list of documents
|
||||
entities: Optional pre-extracted entities (single list or list of lists)
|
||||
relations: Optional pre-extracted relations (single list or list of lists)
|
||||
pipeline_id: Optional pipeline ID for progress tracking
|
||||
**kwargs: Extraction options
|
||||
|
||||
Returns:
|
||||
Union[List[Triplet], List[List[Triplet]]]: Extracted triplets
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
# Handle batch extraction with progress tracking
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="TripletExtractor",
|
||||
message=f"Batch extracting triplets from {len(text)} documents",
|
||||
pipeline_id=pipeline_id,
|
||||
)
|
||||
|
||||
try:
|
||||
results = [None] * len(text)
|
||||
total_items = len(text)
|
||||
total_triplets_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
# Initial progress update
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=0,
|
||||
total=total_items,
|
||||
message=f"Starting batch extraction... 0/{total_items}"
|
||||
)
|
||||
|
||||
from .config import resolve_max_workers
|
||||
max_workers = resolve_max_workers(
|
||||
explicit=kwargs.get("max_workers"),
|
||||
local_config=self.config,
|
||||
methods=self.method,
|
||||
)
|
||||
|
||||
def process_item(idx, item):
|
||||
try:
|
||||
# Prepare arguments for single item
|
||||
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||
|
||||
doc_entities = None
|
||||
if entities and isinstance(entities, list) and idx < len(entities):
|
||||
doc_entities = entities[idx]
|
||||
|
||||
doc_relations = None
|
||||
if relations and isinstance(relations, list) and idx < len(relations):
|
||||
doc_relations = relations[idx]
|
||||
|
||||
# Extract
|
||||
current_triplets = self.extract_triplets(
|
||||
doc_text,
|
||||
entities=doc_entities,
|
||||
relations=doc_relations,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# Add provenance metadata
|
||||
for triplet in current_triplets:
|
||||
if triplet.metadata is None:
|
||||
triplet.metadata = {}
|
||||
triplet.metadata["batch_index"] = idx
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
triplet.metadata["document_id"] = item["id"]
|
||||
|
||||
return idx, current_triplets
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to process item {idx}: {e}")
|
||||
return idx, []
|
||||
|
||||
if max_workers > 1:
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit tasks
|
||||
future_to_idx = {
|
||||
executor.submit(process_item, idx, item): idx
|
||||
for idx, item in enumerate(text)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
idx, triplets = future.result()
|
||||
results[idx] = triplets
|
||||
total_triplets_count += len(triplets)
|
||||
processed_count += 1
|
||||
|
||||
if processed_count % update_interval == 0 or processed_count == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining}) - Extracted {total_triplets_count} triplets"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
for idx, item in enumerate(text):
|
||||
_, triplets = process_item(idx, item)
|
||||
results[idx] = triplets
|
||||
total_triplets_count += len(triplets)
|
||||
processed_count += 1
|
||||
|
||||
if processed_count % update_interval == 0 or processed_count == total_items:
|
||||
remaining = total_items - processed_count
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=processed_count,
|
||||
total=total_items,
|
||||
message=f"Processing... {processed_count}/{total_items} (remaining: {remaining}) - Extracted {total_triplets_count} triplets"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_triplets_count} triplets.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
else:
|
||||
# Single item
|
||||
return self.extract_triplets(text, entities=entities, relations=relations, **kwargs)
|
||||
|
||||
def extract_triplets(
|
||||
self,
|
||||
text: str,
|
||||
entities: Optional[List[Entity]] = None,
|
||||
relations: Optional[List[Relation]] = None,
|
||||
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||
pipeline_id: Optional[str] = None,
|
||||
**options,
|
||||
) -> List[Triplet]:
|
||||
) -> Union[List[Triplet], List[List[Triplet]]]:
|
||||
"""
|
||||
Extract RDF triplets from text.
|
||||
|
||||
@@ -163,11 +331,29 @@ class TripletExtractor:
|
||||
text: Input text
|
||||
entities: Pre-extracted entities (optional)
|
||||
relations: Pre-extracted relations (optional)
|
||||
pipeline_id: Optional pipeline ID for progress tracking (batch mode)
|
||||
**options: Extraction options
|
||||
|
||||
Returns:
|
||||
list: List of extracted triplets
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
entities_batch = entities
|
||||
if entities is not None and isinstance(entities, list) and (not entities or all(isinstance(e, Entity) for e in entities)):
|
||||
entities_batch = [entities for _ in range(len(text))] if entities else [[] for _ in range(len(text))]
|
||||
|
||||
relations_batch = relations
|
||||
if relations is not None and isinstance(relations, list) and (not relations or all(isinstance(r, Relation) for r in relations)):
|
||||
relations_batch = [relations for _ in range(len(text))] if relations else [[] for _ in range(len(text))]
|
||||
|
||||
return self.extract(
|
||||
text,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
pipeline_id=pipeline_id,
|
||||
**options,
|
||||
)
|
||||
|
||||
from .methods import get_triplet_method
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
@@ -185,22 +371,46 @@ class TripletExtractor:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Extracting entities..."
|
||||
)
|
||||
ner = NERExtractor(**self.config.get("ner", {}))
|
||||
entities = ner.extract_entities(text)
|
||||
if self._ner_extractor is None:
|
||||
ner_config = self.config.get("ner", {})
|
||||
if "ner_method" in self.config:
|
||||
ner_config = {**ner_config, "method": self.config["ner_method"]}
|
||||
self._ner_extractor = NERExtractor(
|
||||
**ner_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation", "validator", "serializer", "quality"]
|
||||
},
|
||||
)
|
||||
entities = self._ner_extractor.extract_entities(text)
|
||||
|
||||
# Extract relations if not provided
|
||||
if relations is None:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Extracting relations..."
|
||||
)
|
||||
rel_extractor = RelationExtractor(**self.config.get("relation", {}))
|
||||
relations = rel_extractor.extract_relations(text, entities)
|
||||
if self._relation_extractor is None:
|
||||
rel_config = self.config.get("relation", {})
|
||||
if "relation_method" in self.config:
|
||||
rel_config = {**rel_config, "method": self.config["relation_method"]}
|
||||
self._relation_extractor = RelationExtractor(
|
||||
**rel_config,
|
||||
**{
|
||||
k: v
|
||||
for k, v in self.config.items()
|
||||
if k not in ["ner", "relation", "validator", "serializer", "quality"]
|
||||
},
|
||||
)
|
||||
relations = self._relation_extractor.extract_relations(text, entities)
|
||||
|
||||
# Use method-based extraction
|
||||
methods = options.get("method", self.method)
|
||||
if isinstance(methods, str):
|
||||
methods = [methods]
|
||||
|
||||
triplet_types = options.get("triplet_types", self.triplet_types)
|
||||
|
||||
# Merge config with options
|
||||
all_options = {**self.config, **options}
|
||||
|
||||
@@ -234,6 +444,11 @@ class TripletExtractor:
|
||||
|
||||
# Prepare method-specific options
|
||||
method_options = all_options.copy()
|
||||
|
||||
# Pass triplet_types to all methods
|
||||
if triplet_types:
|
||||
method_options["triplet_types"] = triplet_types
|
||||
|
||||
if method_name == "huggingface":
|
||||
method_options["model"] = all_options.get(
|
||||
"huggingface_model", all_options.get("model")
|
||||
@@ -246,18 +461,28 @@ class TripletExtractor:
|
||||
method_options["model"] = all_options.get(
|
||||
"llm_model", all_options.get("model")
|
||||
)
|
||||
# Pass api_key if provided (needed for all providers)
|
||||
if "api_key" in all_options:
|
||||
method_options["api_key"] = all_options["api_key"]
|
||||
elif "api_key" not in method_options:
|
||||
# Try to get from environment as fallback
|
||||
# Ensure api_key is populated: check explicitly provided or fallback to env
|
||||
current_key = method_options.get("api_key")
|
||||
if not current_key:
|
||||
# Not found or empty/None, try environment
|
||||
import os
|
||||
provider = method_options.get("provider", "openai")
|
||||
env_key = f"{provider.upper()}_API_KEY"
|
||||
provider_name = method_options.get("provider", "openai")
|
||||
env_key = f"{provider_name.upper()}_API_KEY"
|
||||
api_key = os.getenv(env_key)
|
||||
if api_key:
|
||||
method_options["api_key"] = api_key
|
||||
|
||||
# Print progress if verbose mode is enabled (only for LLM method to avoid spam)
|
||||
verbose_mode = options.get("verbose", False) or self.config.get("verbose", False)
|
||||
if verbose_mode and method_name == "llm":
|
||||
import sys
|
||||
print(f" [TripletExtractor] Processing with {method_name}...", flush=True, file=sys.stdout)
|
||||
if "api_key" in method_options:
|
||||
masked = method_options["api_key"][:4] + "..." if method_options["api_key"] else "None"
|
||||
print(f" [TripletExtractor Debug] api_key present: {masked}", flush=True, file=sys.stdout)
|
||||
else:
|
||||
print(f" [TripletExtractor Debug] api_key NOT present", flush=True, file=sys.stdout)
|
||||
|
||||
triplets = method_func(
|
||||
text,
|
||||
entities=entities,
|
||||
@@ -265,6 +490,25 @@ class TripletExtractor:
|
||||
**method_options,
|
||||
)
|
||||
|
||||
# Print result count if verbose (only for LLM method)
|
||||
if verbose_mode and method_name == "llm" and len(triplets) > 0:
|
||||
import sys
|
||||
print(f" [TripletExtractor] Extracted {len(triplets)} triplets", flush=True, file=sys.stdout)
|
||||
|
||||
# Apply weighted scoring if triplet_types are provided
|
||||
if triplet_types:
|
||||
try:
|
||||
from .methods import calculate_weighted_confidence
|
||||
for t in triplets:
|
||||
t.confidence = calculate_weighted_confidence(
|
||||
item_type=t.predicate,
|
||||
original_confidence=t.confidence,
|
||||
valid_types=triplet_types,
|
||||
item_text=t.predicate # For triplets, predicate is the key text
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Filter by confidence
|
||||
min_conf = options.get("min_confidence", self.min_confidence)
|
||||
filtered = [t for t in triplets if t.confidence >= min_conf]
|
||||
@@ -286,6 +530,12 @@ class TripletExtractor:
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Method {method_name} failed: {e}")
|
||||
verbose_mode = options.get("verbose", False) or self.config.get("verbose", False)
|
||||
if verbose_mode:
|
||||
import sys
|
||||
print(f" [TripletExtractor] ERROR: Method {method_name} failed: {e}", flush=True, file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
continue
|
||||
|
||||
# Use first successful method or fallback to relation conversion
|
||||
@@ -293,20 +543,29 @@ class TripletExtractor:
|
||||
triplets = all_triplets[0][1]
|
||||
else:
|
||||
# Fallback: Convert relations to triplets
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id,
|
||||
message=f"Converting {len(relations)} relations to triplets...",
|
||||
)
|
||||
triplets = []
|
||||
for relation in relations:
|
||||
triplet = Triplet(
|
||||
subject=self._format_uri(relation.subject.text),
|
||||
predicate=self._format_uri(relation.predicate),
|
||||
object=self._format_uri(relation.object.text),
|
||||
confidence=relation.confidence,
|
||||
metadata={"context": relation.context, **relation.metadata},
|
||||
if relations:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id,
|
||||
message=f"Converting {len(relations)} relations to triplets...",
|
||||
)
|
||||
triplets.append(triplet)
|
||||
triplets = []
|
||||
for relation in relations:
|
||||
triplet = Triplet(
|
||||
subject=self._format_uri(relation.subject.text),
|
||||
predicate=self._format_uri(relation.predicate),
|
||||
object=self._format_uri(relation.object.text),
|
||||
confidence=relation.confidence,
|
||||
metadata={"context": relation.context, **relation.metadata},
|
||||
)
|
||||
triplets.append(triplet)
|
||||
else:
|
||||
# Last resort: Try rule-based extraction if no relations exist
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id,
|
||||
message="No relations found. Trying rule-based triplet extraction...",
|
||||
)
|
||||
method_func = get_triplet_method("rules")
|
||||
triplets = method_func(text, entities=entities, relations=[], **all_options)
|
||||
|
||||
# Validate triplets
|
||||
if options.get("validate", self._should_validate):
|
||||
@@ -382,7 +641,65 @@ class TripletExtractor:
|
||||
Returns:
|
||||
list: List of triplet lists for each text
|
||||
"""
|
||||
return [self.extract_triplets(text, **options) for text in texts]
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="semantic_extract",
|
||||
submodule="TripletExtractor",
|
||||
message=f"Batch extracting triplets from {len(texts)} documents",
|
||||
)
|
||||
|
||||
results = []
|
||||
total_triplets_count = 0
|
||||
total_items = len(texts)
|
||||
|
||||
try:
|
||||
# Determine update interval
|
||||
if total_items <= 10:
|
||||
update_interval = 1
|
||||
else:
|
||||
update_interval = max(1, min(10, total_items // 100))
|
||||
|
||||
for idx, text in enumerate(texts, 1):
|
||||
# Extract triplets
|
||||
triplets = self.extract_triplets(text, **options)
|
||||
|
||||
# Add provenance metadata
|
||||
for triplet in triplets:
|
||||
if triplet.metadata is None:
|
||||
triplet.metadata = {}
|
||||
triplet.metadata["batch_index"] = idx - 1
|
||||
|
||||
results.append(triplets)
|
||||
total_triplets_count += len(triplets)
|
||||
|
||||
# Update progress
|
||||
should_update = (
|
||||
idx % update_interval == 0 or
|
||||
idx == total_items or
|
||||
idx == 1 or
|
||||
total_items <= 10
|
||||
)
|
||||
|
||||
if should_update:
|
||||
remaining = total_items - idx
|
||||
self.progress_tracker.update_progress(
|
||||
tracking_id,
|
||||
processed=idx,
|
||||
total=total_items,
|
||||
message=f"Processing documents... {idx}/{total_items} (remaining: {remaining}) - Extracted {total_triplets_count} triplets so far"
|
||||
)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_triplets_count} triplets.",
|
||||
)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class TripletValidator:
|
||||
@@ -428,27 +745,6 @@ class TripletValidator:
|
||||
"""
|
||||
return [t for t in triplets if self.validate_triplet(t, **criteria)]
|
||||
|
||||
def check_triplet_consistency(self, triplets: List[Triplet]) -> Dict[str, Any]:
|
||||
"""
|
||||
Check consistency among triplets.
|
||||
|
||||
Args:
|
||||
triplets: List of triplets
|
||||
|
||||
Returns:
|
||||
dict: Consistency report
|
||||
"""
|
||||
issues = []
|
||||
|
||||
# Check for contradictory triplets
|
||||
# (simplified - would need domain knowledge for full implementation)
|
||||
|
||||
return {
|
||||
"total_triplets": len(triplets),
|
||||
"issues": issues,
|
||||
"consistent": len(issues) == 0,
|
||||
}
|
||||
|
||||
|
||||
class RDFSerializer:
|
||||
"""RDF serialization handler."""
|
||||
|
||||
@@ -34,6 +34,7 @@ License: MIT
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
@@ -46,6 +47,13 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from .logging import get_logger
|
||||
|
||||
DISABLE_JUPYTER_PROGRESS = os.getenv("SEMANTICA_DISABLE_JUPYTER_PROGRESS", "").strip().lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
)
|
||||
|
||||
# Try to import IPython for Jupyter support
|
||||
try:
|
||||
from IPython import get_ipython
|
||||
@@ -1017,6 +1025,7 @@ class ProgressTracker:
|
||||
|
||||
# Detect environment - will be checked dynamically
|
||||
self.is_jupyter = self._detect_jupyter()
|
||||
self.disable_jupyter_progress = DISABLE_JUPYTER_PROGRESS
|
||||
|
||||
# Create displays
|
||||
self.displays: List[ProgressDisplay] = []
|
||||
@@ -1024,7 +1033,7 @@ class ProgressTracker:
|
||||
# Always try Jupyter first if available, fallback to console
|
||||
if IPYTHON_AVAILABLE:
|
||||
# Try to detect Jupyter - if available, use it
|
||||
if self.is_jupyter:
|
||||
if self.is_jupyter and not self.disable_jupyter_progress:
|
||||
self.displays.append(JupyterProgressDisplay(use_emoji=use_emoji))
|
||||
# Also add console as fallback for immediate feedback
|
||||
self.displays.append(
|
||||
@@ -1213,7 +1222,11 @@ class ProgressTracker:
|
||||
if IPYTHON_AVAILABLE and not self.is_jupyter:
|
||||
self.is_jupyter = self._detect_jupyter()
|
||||
# If Jupyter is now detected and we don't have a Jupyter display, add it
|
||||
if self.is_jupyter and not any(isinstance(d, JupyterProgressDisplay) for d in self.displays):
|
||||
if (
|
||||
self.is_jupyter
|
||||
and not self.disable_jupyter_progress
|
||||
and not any(isinstance(d, JupyterProgressDisplay) for d in self.displays)
|
||||
):
|
||||
# Insert Jupyter display at the beginning for priority
|
||||
self.displays.insert(0, JupyterProgressDisplay(use_emoji=self.use_emoji))
|
||||
|
||||
@@ -1341,7 +1354,11 @@ class ProgressTracker:
|
||||
if IPYTHON_AVAILABLE and not self.is_jupyter:
|
||||
self.is_jupyter = self._detect_jupyter()
|
||||
# If Jupyter is now detected and we don't have a Jupyter display, add it
|
||||
if self.is_jupyter and not any(isinstance(d, JupyterProgressDisplay) for d in self.displays):
|
||||
if (
|
||||
self.is_jupyter
|
||||
and not self.disable_jupyter_progress
|
||||
and not any(isinstance(d, JupyterProgressDisplay) for d in self.displays)
|
||||
):
|
||||
# Insert Jupyter display at the beginning for priority
|
||||
self.displays.insert(0, JupyterProgressDisplay(use_emoji=self.use_emoji))
|
||||
|
||||
@@ -1526,7 +1543,11 @@ def get_progress_tracker() -> ProgressTracker:
|
||||
if IPYTHON_AVAILABLE and not _global_tracker.is_jupyter:
|
||||
_global_tracker.is_jupyter = _global_tracker._detect_jupyter()
|
||||
# If Jupyter is now detected and we don't have a Jupyter display, add it
|
||||
if _global_tracker.is_jupyter and not any(isinstance(d, JupyterProgressDisplay) for d in _global_tracker.displays):
|
||||
if (
|
||||
_global_tracker.is_jupyter
|
||||
and not _global_tracker.disable_jupyter_progress
|
||||
and not any(isinstance(d, JupyterProgressDisplay) for d in _global_tracker.displays)
|
||||
):
|
||||
# Insert Jupyter display at the beginning for priority
|
||||
_global_tracker.displays.insert(0, JupyterProgressDisplay(use_emoji=_global_tracker.use_emoji))
|
||||
|
||||
|
||||
@@ -38,6 +38,7 @@ License: MIT
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
import concurrent.futures
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -61,7 +62,7 @@ class VectorStore:
|
||||
|
||||
SUPPORTED_BACKENDS = {"faiss", "weaviate", "qdrant", "milvus", "inmemory"}
|
||||
|
||||
def __init__(self, backend="faiss", config=None, **kwargs):
|
||||
def __init__(self, backend="faiss", config=None, max_workers: int = 6, **kwargs):
|
||||
"""Initialize vector store."""
|
||||
if backend.lower() not in self.SUPPORTED_BACKENDS:
|
||||
raise ValueError(
|
||||
@@ -72,6 +73,7 @@ class VectorStore:
|
||||
self.logger = get_logger("vector_store")
|
||||
self.config = config or {}
|
||||
self.config.update(kwargs)
|
||||
self.max_workers = max_workers
|
||||
self.progress_tracker = get_progress_tracker()
|
||||
# Ensure progress tracker is enabled
|
||||
if not self.progress_tracker.enabled:
|
||||
@@ -127,6 +129,141 @@ class VectorStore:
|
||||
self.logger.warning("Using random fallback embedding")
|
||||
return np.random.rand(self.dimension).astype(np.float32)
|
||||
|
||||
def embed_batch(self, texts: List[str]) -> List[np.ndarray]:
|
||||
"""
|
||||
Generate embeddings for a list of texts using the internal embedder.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed
|
||||
|
||||
Returns:
|
||||
List of numpy arrays
|
||||
"""
|
||||
if self.embedder:
|
||||
try:
|
||||
# generate_embeddings handles list input
|
||||
embeddings = self.embedder.generate_embeddings(texts)
|
||||
# Ensure it returns a list of arrays (it returns 2D array or list)
|
||||
if isinstance(embeddings, np.ndarray):
|
||||
return list(embeddings)
|
||||
return embeddings
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Batch embedding generation failed: {e}")
|
||||
|
||||
# Fallback
|
||||
self.logger.warning("Using random fallback embeddings for batch")
|
||||
return [np.random.rand(self.dimension).astype(np.float32) for _ in texts]
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
documents: List[str],
|
||||
metadata: Optional[List[Dict[str, Any]]] = None,
|
||||
batch_size: int = 32,
|
||||
parallel: bool = True,
|
||||
**options,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add multiple documents to the store with parallel embedding generation.
|
||||
|
||||
Args:
|
||||
documents: List of document texts
|
||||
metadata: List of metadata dictionaries
|
||||
batch_size: Number of documents to process in one batch
|
||||
parallel: Whether to use parallel processing for embeddings
|
||||
**options: Additional options
|
||||
|
||||
Returns:
|
||||
List[str]: Vector IDs
|
||||
"""
|
||||
if not documents:
|
||||
return []
|
||||
|
||||
num_docs = len(documents)
|
||||
metadata = metadata or [{} for _ in range(num_docs)]
|
||||
|
||||
if len(metadata) != num_docs:
|
||||
raise ValueError("Metadata list length must match documents length")
|
||||
|
||||
all_vectors = [None] * num_docs
|
||||
|
||||
# Helper for processing a batch
|
||||
def process_batch(start_idx: int, end_idx: int):
|
||||
batch_texts = documents[start_idx:end_idx]
|
||||
batch_embeddings = self.embed_batch(batch_texts)
|
||||
return start_idx, batch_embeddings
|
||||
|
||||
# Calculate batches
|
||||
batches = []
|
||||
for i in range(0, num_docs, batch_size):
|
||||
batches.append((i, min(i + batch_size, num_docs)))
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
module="vector_store",
|
||||
submodule="VectorStore",
|
||||
message=f"Processing {num_docs} documents (parallel={parallel})",
|
||||
)
|
||||
|
||||
try:
|
||||
if parallel and self.max_workers > 1:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message=f"Embedding with {self.max_workers} workers..."
|
||||
)
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(process_batch, start, end)
|
||||
for start, end in batches
|
||||
]
|
||||
|
||||
completed = 0
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
start_idx, embeddings = future.result()
|
||||
# Place results in correct order
|
||||
for i, emb in enumerate(embeddings):
|
||||
all_vectors[start_idx + i] = emb
|
||||
|
||||
completed += 1
|
||||
if completed % 5 == 0: # Update progress periodically
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id,
|
||||
message=f"Embedded batch {completed}/{len(batches)}"
|
||||
)
|
||||
else:
|
||||
# Sequential processing
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Embedding sequentially..."
|
||||
)
|
||||
for i, (start, end) in enumerate(batches):
|
||||
_, embeddings = process_batch(start, end)
|
||||
for j, emb in enumerate(embeddings):
|
||||
all_vectors[start + j] = emb
|
||||
|
||||
if i % 5 == 0:
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id,
|
||||
message=f"Embedded batch {i+1}/{len(batches)}"
|
||||
)
|
||||
|
||||
# Verify all embeddings generated
|
||||
if any(v is None for v in all_vectors):
|
||||
raise ProcessingError("Failed to generate all embeddings")
|
||||
|
||||
# Store all vectors in one go
|
||||
self.progress_tracker.update_tracking(tracking_id, message="Storing vectors...")
|
||||
vector_ids = self.store_vectors(all_vectors, metadata=metadata, **options)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Added {len(vector_ids)} documents",
|
||||
)
|
||||
return vector_ids
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
def store(
|
||||
self,
|
||||
vectors: List[np.ndarray],
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
import os
|
||||
import unittest
|
||||
import time
|
||||
|
||||
from semantica.semantic_extract import NERExtractor, RelationExtractor
|
||||
|
||||
# Use environment variable for API key
|
||||
_GROQ_KEY = os.getenv("GROQ_API_KEY") or os.getenv("GROQ_TEST_API_KEY")
|
||||
@unittest.skipUnless(_GROQ_KEY, "Groq key not set; skipping live integration test")
|
||||
class TestGroqRelationsIntegration(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.api_key = _GROQ_KEY
|
||||
self.model = "llama-3.1-8b-instant"
|
||||
# Short, unambiguous finance snippet
|
||||
self.text_short = (
|
||||
"Apple reported revenue of $4.4 billion in Q1 2024 and provided guidance for FY 2025."
|
||||
)
|
||||
# Longer text to exercise chunking and ensure no hang
|
||||
self.text_long = (
|
||||
"Apple reported revenue of $4.4 billion in Q1 2024. "
|
||||
"The company also reported growth of 12% year-over-year and provided guidance for FY 2025. "
|
||||
"Microsoft reported revenue of $6.1 billion in Q2 2024 and expects sequential growth. "
|
||||
"NVIDIA reported record revenue in 2024 Q1 and guided for higher revenue in Q2 2024. "
|
||||
) * 20 # expand length
|
||||
|
||||
def _extract_entities(self, text):
|
||||
ner = NERExtractor(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model=self.model,
|
||||
api_key=self.api_key,
|
||||
temperature=0.0,
|
||||
)
|
||||
entities = ner.extract_entities(text, entity_types=["ORGANIZATION", "MONEY", "DATE", "EVENT", "PERCENT"])
|
||||
self.assertIsInstance(entities, list)
|
||||
return entities
|
||||
|
||||
def test_relations_short_text(self):
|
||||
entities = self._extract_entities(self.text_short)
|
||||
self.assertGreater(len(entities), 0, "NER should extract entities for short text")
|
||||
|
||||
relation_extractor = RelationExtractor(
|
||||
method="llm",
|
||||
relation_types=[
|
||||
"HAS_REVENUE",
|
||||
"HAS_GROWTH",
|
||||
"PROVIDES_GUIDANCE",
|
||||
"IN_QUARTER",
|
||||
"FOR_PERIOD",
|
||||
"RELATED_TO",
|
||||
],
|
||||
provider="groq",
|
||||
llm_model=self.model,
|
||||
api_key=self.api_key,
|
||||
temperature=0.0,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
relations = relation_extractor.extract_relations(text=self.text_short, entities=entities)
|
||||
elapsed = time.time() - start
|
||||
|
||||
self.assertIsInstance(relations, list)
|
||||
# Ensure call completes reasonably fast (network dependent; allow generous bound)
|
||||
self.assertLess(elapsed, 60, f"Extraction took too long: {elapsed:.2f}s")
|
||||
# Do not strictly assert >0 as model output may vary, but log for diagnostics
|
||||
if relations:
|
||||
sample = relations[0]
|
||||
self.assertTrue(hasattr(sample, "subject") and hasattr(sample, "predicate") and hasattr(sample, "object"))
|
||||
|
||||
def test_relations_long_text_chunking(self):
|
||||
entities = self._extract_entities(self.text_long)
|
||||
self.assertGreater(len(entities), 0, "NER should extract entities for long text")
|
||||
|
||||
relation_extractor = RelationExtractor(
|
||||
method="llm",
|
||||
relation_types=["RELATED_TO", "HAS_REVENUE", "IN_QUARTER"],
|
||||
provider="groq",
|
||||
llm_model=self.model,
|
||||
api_key=self.api_key,
|
||||
temperature=0.0,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
relations = relation_extractor.extract_relations(text=self.text_long, entities=entities)
|
||||
elapsed = time.time() - start
|
||||
|
||||
self.assertIsInstance(relations, list)
|
||||
# Ensure completion (chunked path) and no hang
|
||||
self.assertLess(elapsed, 120, f"Chunked extraction took too long: {elapsed:.2f}s")
|
||||
if relations:
|
||||
for r in relations[:3]:
|
||||
self.assertTrue(hasattr(r, "subject") and hasattr(r, "predicate") and hasattr(r, "object"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,232 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from semantica.kg.graph_builder import GraphBuilder
|
||||
|
||||
|
||||
class TestGraphBuilderExternal(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.mock_tracker_patcher = patch("semantica.utils.progress_tracker.get_progress_tracker")
|
||||
self.mock_get_tracker = self.mock_tracker_patcher.start()
|
||||
self.mock_tracker = MagicMock()
|
||||
self.mock_get_tracker.return_value = self.mock_tracker
|
||||
|
||||
self.mock_resolver_patcher = patch("semantica.kg.entity_resolver.EntityResolver")
|
||||
self.mock_resolver_cls = self.mock_resolver_patcher.start()
|
||||
|
||||
self.mock_conflict_patcher = patch("semantica.conflicts.conflict_detector.ConflictDetector")
|
||||
self.mock_conflict_cls = self.mock_conflict_patcher.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.mock_tracker_patcher.stop()
|
||||
self.mock_resolver_patcher.stop()
|
||||
self.mock_conflict_patcher.stop()
|
||||
|
||||
def test_single_source_dict_with_source_id_target_id(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
entities = [
|
||||
{"id": "drug:1", "name": "Aspirin", "type": "Drug"},
|
||||
{"id": "disease:1", "name": "Myocardial infarction", "type": "Disease"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "drug:1", "target_id": "disease:1", "type": "TREATS"},
|
||||
]
|
||||
|
||||
source = {"entities": entities, "relationships": relationships}
|
||||
|
||||
kg = builder.build(source)
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 2)
|
||||
self.assertEqual(len(kg["relationships"]), 1)
|
||||
rel = kg["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "drug:1")
|
||||
self.assertEqual(rel.get("target"), "disease:1")
|
||||
self.assertEqual(kg["metadata"]["num_relationships"], 1)
|
||||
|
||||
def test_sources_list_merge_with_external_relationships(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
source1 = {
|
||||
"entities": [{"id": "1", "name": "A"}],
|
||||
"relationships": [{"source_id": "1", "target_id": "2", "type": "REL_1"}],
|
||||
}
|
||||
source2 = {
|
||||
"entities": [{"id": "2", "name": "B"}],
|
||||
"relationships": [{"source_id": "2", "target_id": "1", "type": "REL_2"}],
|
||||
}
|
||||
|
||||
kg = builder.build([source1, source2])
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 2)
|
||||
self.assertEqual(len(kg["relationships"]), 2)
|
||||
sources = {r["source"] for r in kg["relationships"]}
|
||||
targets = {r["target"] for r in kg["relationships"]}
|
||||
self.assertEqual(sources, {"1", "2"})
|
||||
self.assertEqual(targets, {"1", "2"})
|
||||
|
||||
def test_build_with_explicit_relationships_argument_external_ids(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
entities = [
|
||||
{"id": "1", "name": "A"},
|
||||
{"id": "2", "name": "B"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "1", "target_id": "2", "type": "REL"},
|
||||
]
|
||||
|
||||
kg = builder.build(entities, relationships=relationships)
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 2)
|
||||
self.assertEqual(len(kg["relationships"]), 1)
|
||||
rel = kg["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "1")
|
||||
self.assertEqual(rel.get("target"), "2")
|
||||
|
||||
def test_build_single_source_external_graph(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
source = {
|
||||
"entities": [{"id": "1", "name": "A"}],
|
||||
"relationships": [{"source_id": "1", "target_id": "1", "type": "SELF"}],
|
||||
}
|
||||
|
||||
kg = builder.build_single_source(source)
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 1)
|
||||
self.assertEqual(len(kg["relationships"]), 1)
|
||||
rel = kg["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "1")
|
||||
self.assertEqual(rel.get("target"), "1")
|
||||
|
||||
def test_relationship_key_variants_normalized(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
entities = [
|
||||
{"id": "1", "name": "A"},
|
||||
{"id": "2", "name": "B"},
|
||||
{"id": "3", "name": "C"},
|
||||
{"id": "4", "name": "D"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "1", "target_id": "2", "type": "R1"},
|
||||
{"source": "2", "target": "3", "type": "R2"},
|
||||
{"subject": "3", "object": "4", "type": "R3"},
|
||||
]
|
||||
|
||||
kg = builder.build({"entities": entities, "relationships": relationships})
|
||||
|
||||
self.assertEqual(len(kg["relationships"]), 3)
|
||||
ids = {(r["source"], r["target"]) for r in kg["relationships"]}
|
||||
self.assertIn(("1", "2"), ids)
|
||||
self.assertIn(("2", "3"), ids)
|
||||
self.assertIn(("3", "4"), ids)
|
||||
|
||||
def test_warning_when_all_relationships_dropped(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
source = {
|
||||
"entities": [],
|
||||
"relationships": [{"foo": "x"}, {"bar": "y"}],
|
||||
}
|
||||
|
||||
with patch.object(builder.logger, "warning") as mock_warning:
|
||||
kg = builder.build(source)
|
||||
|
||||
self.assertEqual(len(kg["relationships"]), 0)
|
||||
mock_warning.assert_called()
|
||||
args, _ = mock_warning.call_args
|
||||
self.assertIn("All relationships were dropped", args[0])
|
||||
|
||||
def test_no_warning_when_some_relationships_kept(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
source = {
|
||||
"entities": [{"id": "1"}, {"id": "2"}],
|
||||
"relationships": [
|
||||
{"source_id": "1", "target_id": "2", "type": "REL"},
|
||||
{"foo": "x"},
|
||||
],
|
||||
}
|
||||
|
||||
with patch.object(builder.logger, "warning") as mock_warning:
|
||||
kg = builder.build(source)
|
||||
|
||||
self.assertEqual(len(kg["relationships"]), 2)
|
||||
mock_warning.assert_not_called()
|
||||
|
||||
def test_issue_208_minimal_reproduction_shape(self):
|
||||
builder = GraphBuilder(
|
||||
merge_entities=False,
|
||||
entity_resolution_strategy="none",
|
||||
resolve_conflicts=False,
|
||||
)
|
||||
|
||||
entities = [
|
||||
{"id": "e1", "name": "Entity 1"},
|
||||
{"id": "e2", "name": "Entity 2"},
|
||||
{"id": "e3", "name": "Entity 3"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "e1", "target_id": "e2", "type": "REL_1"},
|
||||
{"source_id": "e2", "target_id": "e3", "type": "REL_2"},
|
||||
]
|
||||
|
||||
entity_ids = {e["id"] for e in entities}
|
||||
for r in relationships:
|
||||
self.assertIn(r["source_id"], entity_ids)
|
||||
self.assertIn(r["target_id"], entity_ids)
|
||||
|
||||
kg = builder.build(
|
||||
sources=[{"entities": entities, "relationships": relationships}],
|
||||
merge_entities=False,
|
||||
)
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 3)
|
||||
self.assertEqual(len(kg["relationships"]), 2)
|
||||
pairs = {(r["source"], r["target"]) for r in kg["relationships"]}
|
||||
self.assertIn(("e1", "e2"), pairs)
|
||||
self.assertIn(("e2", "e3"), pairs)
|
||||
|
||||
def test_issue_206_earnings_call_shape(self):
|
||||
builder = GraphBuilder(
|
||||
merge_entities=False,
|
||||
entity_resolution_strategy="none",
|
||||
resolve_conflicts=False,
|
||||
)
|
||||
|
||||
entities = [
|
||||
{
|
||||
"id": "entity_446_MDA Space Ltd.",
|
||||
"name": "MDA Space Ltd.",
|
||||
"type": "ORGANIZATION",
|
||||
},
|
||||
{
|
||||
"id": "entity_500_$409.8 million",
|
||||
"name": "$409.8 million",
|
||||
"type": "MONEY",
|
||||
},
|
||||
]
|
||||
|
||||
relationships = [
|
||||
{
|
||||
"id": None,
|
||||
"source_id": "MDA Space Ltd.",
|
||||
"target_id": "$409.8 million",
|
||||
"type": "HAS_REVENUE",
|
||||
"confidence": 0.975,
|
||||
"metadata": {},
|
||||
}
|
||||
]
|
||||
|
||||
kg = builder.build(
|
||||
sources=[{"entities": entities, "relationships": relationships}],
|
||||
merge_entities=False,
|
||||
)
|
||||
|
||||
self.assertEqual(len(kg["entities"]), 2)
|
||||
self.assertEqual(len(kg["relationships"]), 1)
|
||||
rel = kg["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "MDA Space Ltd.")
|
||||
self.assertEqual(rel.get("target"), "$409.8 million")
|
||||
@@ -91,6 +91,30 @@ class TestGraphBuilder(unittest.TestCase):
|
||||
graph2 = builder.build(source_list)
|
||||
self.assertEqual(len(graph2["entities"]), 2)
|
||||
|
||||
def test_build_with_external_relationship_ids(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
entities = [
|
||||
{"id": "1", "name": "A"},
|
||||
{"id": "2", "name": "B"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "1", "target_id": "2", "type": "rel"},
|
||||
]
|
||||
|
||||
source = {
|
||||
"entities": entities,
|
||||
"relationships": relationships,
|
||||
}
|
||||
|
||||
graph = builder.build(source)
|
||||
|
||||
self.assertEqual(len(graph["entities"]), 2)
|
||||
self.assertEqual(len(graph["relationships"]), 1)
|
||||
rel = graph["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "1")
|
||||
self.assertEqual(rel.get("target"), "2")
|
||||
|
||||
def test_build_with_conflict_resolution(self):
|
||||
"""Test building with conflict resolution enabled"""
|
||||
builder = GraphBuilder(resolve_conflicts=True)
|
||||
@@ -106,6 +130,50 @@ class TestGraphBuilder(unittest.TestCase):
|
||||
self.mock_conflict_cls.return_value.detect_conflicts.assert_called_once()
|
||||
self.mock_conflict_cls.return_value.resolve_conflicts.assert_called_once()
|
||||
|
||||
def test_build_single_source(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
source = {
|
||||
"entities": [{"id": "1", "name": "A"}],
|
||||
"relationships": [{"source_id": "1", "target_id": "1", "type": "self"}],
|
||||
}
|
||||
graph = builder.build_single_source(source)
|
||||
self.assertEqual(len(graph["entities"]), 1)
|
||||
self.assertEqual(len(graph["relationships"]), 1)
|
||||
|
||||
def test_build_with_explicit_relationships_argument(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
|
||||
entities = [
|
||||
{"id": "1", "name": "A"},
|
||||
{"id": "2", "name": "B"},
|
||||
]
|
||||
relationships = [
|
||||
{"source_id": "1", "target_id": "2", "type": "rel"},
|
||||
]
|
||||
|
||||
graph = builder.build(entities, relationships=relationships)
|
||||
|
||||
self.assertEqual(len(graph["entities"]), 2)
|
||||
self.assertEqual(len(graph["relationships"]), 1)
|
||||
rel = graph["relationships"][0]
|
||||
self.assertEqual(rel.get("source"), "1")
|
||||
self.assertEqual(rel.get("target"), "2")
|
||||
|
||||
def test_build_warns_when_all_relationships_dropped(self):
|
||||
builder = GraphBuilder(merge_entities=False, resolve_conflicts=False)
|
||||
source = {
|
||||
"entities": [],
|
||||
"relationships": [{"foo": "x"}, {"bar": "y"}],
|
||||
}
|
||||
|
||||
with patch.object(builder.logger, "warning") as mock_warning:
|
||||
graph = builder.build(source)
|
||||
|
||||
self.assertEqual(len(graph["relationships"]), 0)
|
||||
mock_warning.assert_called()
|
||||
args, _ = mock_warning.call_args
|
||||
self.assertIn("All relationships were dropped", args[0])
|
||||
|
||||
class TestGraphAnalyzer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.mock_tracker_patcher = patch("semantica.kg.graph_analyzer.get_progress_tracker")
|
||||
|
||||
@@ -107,5 +107,25 @@ class TestPipelineModule(unittest.TestCase):
|
||||
|
||||
self.assertEqual(execution_order, ["A", "B", "C"])
|
||||
|
||||
def test_imports_no_circular_dependencies(self):
|
||||
import semantica
|
||||
|
||||
_ = semantica.pipeline
|
||||
|
||||
from semantica.pipeline import PipelineBuilder, PipelineValidator
|
||||
from semantica.deduplication import DuplicateDetector
|
||||
|
||||
builder = PipelineBuilder()
|
||||
builder.add_step("step1", "dummy")
|
||||
pipeline = builder.build("import_test_pipeline")
|
||||
|
||||
validator = PipelineValidator()
|
||||
result = validator.validate_pipeline(pipeline)
|
||||
|
||||
self.assertTrue(result.valid)
|
||||
|
||||
detector = DuplicateDetector()
|
||||
self.assertIsNotNone(detector)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from semantica.semantic_extract.methods import extract_relations_llm, extract_entities_llm, extract_triplets_llm
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
|
||||
class TestMaxTokensPropagation(unittest.TestCase):
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_max_tokens_propagation_relations(self, mock_create_provider):
|
||||
"""Test that max_tokens is passed to generate_typed in extract_relations_llm."""
|
||||
# Setup mock
|
||||
mock_llm = MagicMock()
|
||||
mock_create_provider.return_value = mock_llm
|
||||
mock_llm.is_available.return_value = True
|
||||
|
||||
# Setup return value to avoid pydantic validation errors
|
||||
mock_response = MagicMock()
|
||||
mock_response.relations = []
|
||||
mock_llm.generate_typed.return_value = mock_response
|
||||
|
||||
# Create dummy entities
|
||||
entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
|
||||
|
||||
# Call the function with max_tokens
|
||||
extract_relations_llm(
|
||||
text="some text",
|
||||
entities=entities,
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_tokens=128000
|
||||
)
|
||||
|
||||
# Check if generate_typed was called with max_tokens
|
||||
args, kwargs = mock_llm.generate_typed.call_args
|
||||
|
||||
print(f"Relations Call kwargs: {kwargs}")
|
||||
|
||||
self.assertIn("max_tokens", kwargs)
|
||||
self.assertEqual(kwargs["max_tokens"], 128000)
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_max_tokens_propagation_entities(self, mock_create_provider):
|
||||
"""Test that max_tokens is passed to generate_typed in extract_entities_llm."""
|
||||
# Setup mock
|
||||
mock_llm = MagicMock()
|
||||
mock_create_provider.return_value = mock_llm
|
||||
mock_llm.is_available.return_value = True
|
||||
|
||||
# Setup return value to avoid pydantic validation errors
|
||||
mock_response = MagicMock()
|
||||
mock_response.entities = []
|
||||
mock_llm.generate_typed.return_value = mock_response
|
||||
|
||||
# Call the function with max_tokens
|
||||
extract_entities_llm(
|
||||
text="some text",
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_tokens=128000
|
||||
)
|
||||
|
||||
# Check if generate_typed was called with max_tokens
|
||||
args, kwargs = mock_llm.generate_typed.call_args
|
||||
|
||||
print(f"Entities Call kwargs: {kwargs}")
|
||||
|
||||
self.assertIn("max_tokens", kwargs)
|
||||
self.assertEqual(kwargs["max_tokens"], 128000)
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_max_tokens_propagation_triplets(self, mock_create_provider):
|
||||
"""Test that max_tokens is passed to generate_typed in extract_triplets_llm."""
|
||||
# Setup mock
|
||||
mock_llm = MagicMock()
|
||||
mock_create_provider.return_value = mock_llm
|
||||
mock_llm.is_available.return_value = True
|
||||
|
||||
# Setup return value to avoid pydantic validation errors
|
||||
mock_response = MagicMock()
|
||||
mock_response.triplets = []
|
||||
mock_llm.generate_typed.return_value = mock_response
|
||||
|
||||
# Call the function with max_tokens
|
||||
extract_triplets_llm(
|
||||
text="some text",
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_tokens=128000
|
||||
)
|
||||
|
||||
# Check if generate_typed was called with max_tokens
|
||||
args, kwargs = mock_llm.generate_typed.call_args
|
||||
|
||||
print(f"Triplets Call kwargs: {kwargs}")
|
||||
|
||||
self.assertIn("max_tokens", kwargs)
|
||||
self.assertEqual(kwargs["max_tokens"], 128000)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,125 @@
|
||||
import statistics
|
||||
import time
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict, List
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
from semantica.semantic_extract.event_detector import EventDetector
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.utils.progress_tracker import get_progress_tracker
|
||||
|
||||
|
||||
def _make_documents(n: int) -> List[Dict[str, str]]:
|
||||
base = (
|
||||
"Apple Inc. was founded by Steve Jobs in 1976 and is headquartered in Cupertino, California. "
|
||||
"Microsoft Corporation was founded by Bill Gates and Paul Allen in 1975. "
|
||||
"In 2014, Apple acquired Beats Electronics for $3 billion. "
|
||||
"In 2023, Google announced a partnership with OpenAI to improve search experiences."
|
||||
)
|
||||
return [{"id": f"doc_{i}", "content": f"{base} Document number {i}."} for i in range(n)]
|
||||
|
||||
|
||||
def _median_seconds(fn, repeats: int = 3) -> float:
|
||||
times = []
|
||||
for _ in range(repeats):
|
||||
start = time.perf_counter()
|
||||
fn()
|
||||
times.append(time.perf_counter() - start)
|
||||
return statistics.median(times)
|
||||
|
||||
|
||||
def _bench(label: str, fn, repeats: int = 3) -> dict:
|
||||
fn()
|
||||
seconds = _median_seconds(fn, repeats=repeats)
|
||||
return {"label": label, "seconds": seconds}
|
||||
|
||||
|
||||
def main():
|
||||
progress = get_progress_tracker()
|
||||
progress.displays = []
|
||||
|
||||
docs = _make_documents(80)
|
||||
texts = [d["content"] for d in docs]
|
||||
|
||||
ner = NERExtractor(method="pattern")
|
||||
rel = RelationExtractor(method="pattern")
|
||||
trip = TripletExtractor(method="pattern")
|
||||
events = EventDetector(method="pattern")
|
||||
analyzer = SemanticAnalyzer()
|
||||
net = SemanticNetworkExtractor(ner_method="pattern", relation_method="pattern")
|
||||
|
||||
results = []
|
||||
|
||||
def ner_parallel():
|
||||
ner.extract(texts)
|
||||
|
||||
def ner_seq():
|
||||
ner.extract(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("NER batch (default workers)", ner_parallel))
|
||||
results.append(_bench("NER batch (max_workers=1)", ner_seq))
|
||||
|
||||
entities_batch = ner.extract(texts)
|
||||
|
||||
def rel_parallel():
|
||||
rel.extract(texts, entities_batch)
|
||||
|
||||
def rel_seq():
|
||||
rel.extract(texts, entities_batch, max_workers=1)
|
||||
|
||||
results.append(_bench("Relation batch (default workers)", rel_parallel))
|
||||
results.append(_bench("Relation batch (max_workers=1)", rel_seq))
|
||||
|
||||
def trip_parallel():
|
||||
trip.extract(texts)
|
||||
|
||||
def trip_seq():
|
||||
trip.extract(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Triplet pipeline (default workers)", trip_parallel))
|
||||
results.append(_bench("Triplet pipeline (max_workers=1)", trip_seq))
|
||||
|
||||
def ev_parallel():
|
||||
events.detect_events(texts)
|
||||
|
||||
def ev_seq():
|
||||
events.detect_events(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Event detection (default workers)", ev_parallel))
|
||||
results.append(_bench("Event detection (max_workers=1)", ev_seq))
|
||||
|
||||
def analyzer_parallel():
|
||||
analyzer.analyze(texts)
|
||||
|
||||
def analyzer_seq():
|
||||
analyzer.analyze(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Semantic analysis (default workers)", analyzer_parallel))
|
||||
results.append(_bench("Semantic analysis (max_workers=1)", analyzer_seq))
|
||||
|
||||
def net_parallel():
|
||||
net.extract_network(texts)
|
||||
|
||||
def net_seq():
|
||||
net.extract_network(texts, max_workers=1)
|
||||
|
||||
results.append(_bench("Semantic network (default workers)", net_parallel))
|
||||
results.append(_bench("Semantic network (max_workers=1)", net_seq))
|
||||
|
||||
per_doc = []
|
||||
for row in results:
|
||||
per_doc.append({**row, "ms_per_doc": (row["seconds"] / len(texts)) * 1000.0})
|
||||
|
||||
print(f"Documents: {len(texts)}")
|
||||
for row in per_doc:
|
||||
print(f"{row['label']}: {row['seconds']:.3f}s ({row['ms_per_doc']:.2f} ms/doc)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -7,8 +7,12 @@ import os
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.ner_extractor import Entity as NEREntity
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.semantic_extract.event_detector import EventDetector
|
||||
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.named_entity_recognizer import Entity
|
||||
from semantica.semantic_extract.relation_extractor import Relation
|
||||
|
||||
@@ -51,6 +55,21 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(entities, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_entity_method")
|
||||
def test_ner_extraction_batch_via_extract_entities(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_entities.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = NERExtractor(method="pattern")
|
||||
results = extractor.extract_entities(
|
||||
["Test text 1", "Test text 2"],
|
||||
)
|
||||
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_relation_method")
|
||||
def test_relation_extraction(self, mock_get_method):
|
||||
"""Test relation extraction call"""
|
||||
@@ -65,6 +84,21 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(relations, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_relation_method")
|
||||
def test_relation_extraction_batch_via_extract_relations(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_relations.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = RelationExtractor(method="pattern")
|
||||
texts = ["A knows B", "A knows B"]
|
||||
entities = [NEREntity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)]
|
||||
|
||||
results = extractor.extract_relations(texts, entities)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_triplet_method")
|
||||
def test_triplet_extraction(self, mock_get_method):
|
||||
"""Test triplet extraction call"""
|
||||
@@ -81,5 +115,55 @@ class TestExtractors(unittest.TestCase):
|
||||
self.assertIsInstance(triplets, list)
|
||||
mock_get_method.assert_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.get_triplet_method")
|
||||
def test_triplet_extraction_batch_via_extract_triplets(self, mock_get_method):
|
||||
mock_method = MagicMock()
|
||||
mock_method.extract_triplets.return_value = []
|
||||
mock_get_method.return_value = mock_method
|
||||
|
||||
extractor = TripletExtractor(method="pattern")
|
||||
texts = ["A knows A", "A knows A"]
|
||||
entities_batch = [[NEREntity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)] for _ in texts]
|
||||
relations_batch = [[] for _ in texts]
|
||||
|
||||
results = extractor.extract_triplets(
|
||||
texts,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
)
|
||||
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
def test_event_detector_batch_via_detect_events(self):
|
||||
detector = EventDetector()
|
||||
texts = ["Apple acquired Beats in 2014.", "Google announced a partnership in 2023."]
|
||||
results = detector.detect_events(texts)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, list) for r in results))
|
||||
|
||||
def test_semantic_analyzer_batch_parallel(self):
|
||||
analyzer = SemanticAnalyzer()
|
||||
texts = ["A short sentence.", "Another short sentence."]
|
||||
results = analyzer.analyze(texts)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertTrue(all(isinstance(r, dict) for r in results))
|
||||
|
||||
def test_semantic_network_batch_via_extract_network(self):
|
||||
extractor = SemanticNetworkExtractor()
|
||||
texts = ["A knows B.", "C knows D."]
|
||||
entities_batch = [[] for _ in texts]
|
||||
relations_batch = [[] for _ in texts]
|
||||
results = extractor.extract_network(
|
||||
texts,
|
||||
entities=entities_batch,
|
||||
relations=relations_batch,
|
||||
)
|
||||
self.assertIsInstance(results, list)
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
|
||||
import unittest
|
||||
import time
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.semantic_extract.event_detector import EventDetector
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.methods import _result_cache
|
||||
|
||||
class TestGroqRealWorldPerformance(unittest.TestCase):
|
||||
"""
|
||||
Real-world performance test suite using Groq LLM.
|
||||
Tests parallel processing, caching, and correctness.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.api_key = os.getenv("GROQ_API_KEY")
|
||||
if not cls.api_key:
|
||||
raise unittest.SkipTest("GROQ_API_KEY is not set")
|
||||
cls.metrics_file = os.path.join(os.getcwd(), "groq_metrics.txt")
|
||||
|
||||
# Real-world sample texts (mix of tech, business, and general)
|
||||
cls.sample_texts = [
|
||||
"""Apple Inc. is planning to launch a new AI-powered iPhone in late 2024.
|
||||
CEO Tim Cook announced that the device will feature a neural engine capable of
|
||||
processing 50 trillion operations per second. The company's stock rose 5% following the news.""",
|
||||
|
||||
"""Microsoft Corporation has acquired Activision Blizzard for $68.7 billion.
|
||||
Satya Nadella, Microsoft's Chairman and CEO, stated that this acquisition will
|
||||
accelerate growth in Microsoft's gaming business across mobile, PC, console, and cloud.""",
|
||||
|
||||
"""Elon Musk's SpaceX successfully launched the Starship rocket from Boca Chica, Texas.
|
||||
The mission aims to test new heat shield technology essential for future Mars missions.
|
||||
NASA Administrator Bill Nelson congratulated the team on the achievement.""",
|
||||
|
||||
"""Google DeepMind introduced Gemini, a new multimodal AI model.
|
||||
Sundar Pichai emphasized that Gemini represents a significant leap forward in
|
||||
AI capabilities, outperforming GPT-4 on several benchmarks including MMLU.""",
|
||||
|
||||
"""Amazon Web Services (AWS) announced a partnership with Anthropic to develop
|
||||
reliable and high-performance foundation models. Amazon is investing up to $4 billion
|
||||
in the AI safety startup founded by Dario Amodei."""
|
||||
]
|
||||
|
||||
# Warm up: Ensure modules are loaded
|
||||
print("\n[Setup] Initializing extractors...")
|
||||
cls.extractor = NERExtractor(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model="llama-3.3-70b-versatile",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
cls.relation_extractor = RelationExtractor(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model="llama-3.3-70b-versatile",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
cls.triplet_extractor = TripletExtractor(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model="llama-3.3-70b-versatile",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
cls.event_detector = EventDetector(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model="llama-3.3-70b-versatile",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
cls.network_extractor = SemanticNetworkExtractor(
|
||||
method="llm",
|
||||
provider="groq",
|
||||
llm_model="llama-3.3-70b-versatile",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
|
||||
def setUp(self):
|
||||
# Clear cache before specific performance tests to ensure fair comparison
|
||||
# (Unless testing cache specifically)
|
||||
if _result_cache:
|
||||
_result_cache._caches["entities"].clear()
|
||||
_result_cache._caches["relations"].clear()
|
||||
|
||||
def log_metrics(self, message):
|
||||
print(message)
|
||||
with open(self.metrics_file, "a") as f:
|
||||
f.write(message + "\n")
|
||||
f.flush()
|
||||
|
||||
def test_01_parallel_vs_sequential_performance(self):
|
||||
"""Compare sequential vs parallel extraction speed."""
|
||||
try:
|
||||
self.log_metrics("\n" + "="*60)
|
||||
self.log_metrics("TEST 1: Sequential vs Parallel Processing Performance")
|
||||
self.log_metrics("="*60)
|
||||
|
||||
extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
|
||||
|
||||
# 1. Sequential Run (Max workers = 1)
|
||||
self.log_metrics("\nStarting Sequential Extraction (5 documents)...")
|
||||
start_time = time.time()
|
||||
seq_results = extractor.extract(self.sample_texts, max_workers=1)
|
||||
seq_time = time.time() - start_time
|
||||
self.log_metrics(f"Sequential Time: {seq_time:.4f}s")
|
||||
self.log_metrics(f"Average Latency: {seq_time/len(self.sample_texts):.4f}s per doc")
|
||||
|
||||
# Clear cache to force re-extraction for parallel test
|
||||
_result_cache._caches["entities"].clear()
|
||||
|
||||
# 2. Parallel Run (Max workers = 5)
|
||||
self.log_metrics("\nStarting Parallel Extraction (5 documents, 5 workers)...")
|
||||
start_time = time.time()
|
||||
par_results = extractor.extract(self.sample_texts, max_workers=5)
|
||||
par_time = time.time() - start_time
|
||||
self.log_metrics(f"Parallel Time: {par_time:.4f}s")
|
||||
self.log_metrics(f"Average Latency: {par_time/len(self.sample_texts):.4f}s per doc")
|
||||
|
||||
# Analysis
|
||||
speedup = seq_time / par_time if par_time > 0 else 0
|
||||
self.log_metrics(f"\n>>> Performance Gain: {speedup:.2f}x speedup")
|
||||
self.log_metrics(f">>> Latency Reduction: {(seq_time - par_time):.4f}s total time saved")
|
||||
|
||||
self.assertLess(par_time, seq_time * 1.35, "Parallel processing should not be significantly slower")
|
||||
self.assertEqual(len(seq_results), len(self.sample_texts))
|
||||
self.assertEqual(len(par_results), len(self.sample_texts))
|
||||
except Exception as e:
|
||||
self.log_metrics(f"ERROR in Test 1: {e}")
|
||||
raise
|
||||
|
||||
def test_02_caching_latency_reduction(self):
|
||||
"""Measure latency reduction from caching."""
|
||||
try:
|
||||
self.log_metrics("\n" + "="*60)
|
||||
self.log_metrics("TEST 2: Caching Performance & Latency Reduction")
|
||||
self.log_metrics("="*60)
|
||||
|
||||
extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
|
||||
text = [self.sample_texts[0]]
|
||||
|
||||
# 1. Cold Cache
|
||||
_result_cache._caches["entities"].clear()
|
||||
self.log_metrics("\nCold Cache Request...")
|
||||
start_time = time.time()
|
||||
extractor.extract(text)
|
||||
cold_time = time.time() - start_time
|
||||
self.log_metrics(f"Cold Cache Time: {cold_time:.4f}s")
|
||||
cache_size_after_cold = _result_cache.get_stats()["entities"]["size"]
|
||||
|
||||
# 2. Warm Cache
|
||||
self.log_metrics("\nWarm Cache Request (Identical Query)...")
|
||||
start_time = time.time()
|
||||
extractor.extract(text)
|
||||
warm_time = time.time() - start_time
|
||||
self.log_metrics(f"Warm Cache Time: {warm_time:.6f}s")
|
||||
cache_size_after_warm = _result_cache.get_stats()["entities"]["size"]
|
||||
|
||||
# Analysis
|
||||
reduction = (cold_time - warm_time) / cold_time * 100
|
||||
self.log_metrics(f"\n>>> Latency Reduction: {reduction:.2f}%")
|
||||
|
||||
self.assertLess(warm_time, 1.0, "Warm cache response should be fast (<1.0s)")
|
||||
# self.assertGreater(reduction, 50, "Caching should reduce latency by >50%")
|
||||
if reduction < 30:
|
||||
self.log_metrics(f"WARNING: Caching reduction is low ({reduction:.2f}%)")
|
||||
self.assertGreater(reduction, 20, "Caching should reduce latency by >20%")
|
||||
self.assertGreater(cache_size_after_cold, 0, "Cache should store entity results")
|
||||
self.assertEqual(cache_size_after_warm, cache_size_after_cold, "Warm request should hit the cache")
|
||||
except Exception as e:
|
||||
self.log_metrics(f"ERROR in Test 2: {e}")
|
||||
raise
|
||||
|
||||
def test_03_correctness_and_entity_matching(self):
|
||||
"""Verify extraction correctness and data quality."""
|
||||
try:
|
||||
self.log_metrics("\n" + "="*60)
|
||||
self.log_metrics("TEST 3: Extraction Correctness & Data Quality")
|
||||
self.log_metrics("="*60)
|
||||
|
||||
# Use a specific text with clear entities
|
||||
text = "Satya Nadella is the CEO of Microsoft."
|
||||
|
||||
extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
|
||||
entities = extractor.extract([text])[0] # List of lists
|
||||
|
||||
self.log_metrics(f"\nInput: {text}")
|
||||
self.log_metrics(f"Extracted Entities: {[e.text + '(' + e.label + ')' for e in entities]}")
|
||||
|
||||
# Validation
|
||||
found_person = any(e.label == "PERSON" and "Satya" in e.text for e in entities)
|
||||
found_org = any(e.label == "ORG" and "Microsoft" in e.text for e in entities)
|
||||
|
||||
if not found_org:
|
||||
self.log_metrics("FAILURE: Did not find Microsoft as ORG. Found entities:")
|
||||
for e in entities:
|
||||
self.log_metrics(f" - {e.text}: {e.label}")
|
||||
|
||||
self.assertTrue(found_person, "Failed to extract Satya Nadella as PERSON")
|
||||
self.assertTrue(found_org, "Failed to extract Microsoft as ORG")
|
||||
|
||||
self.log_metrics("\n>>> Correctness Verification: PASS")
|
||||
self.log_metrics(" - Identified PERSON entity")
|
||||
self.log_metrics(" - Identified ORG entity")
|
||||
self.log_metrics(" - Pydantic models validated successfully")
|
||||
except Exception as e:
|
||||
self.log_metrics(f"ERROR in Test 3: {e}")
|
||||
raise
|
||||
|
||||
def test_4_relation_extraction(self):
|
||||
"""Test Relation Extraction capabilities"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 4: Relation Extraction")
|
||||
print("="*60)
|
||||
|
||||
text = self.sample_texts[1] # Microsoft acquisition text
|
||||
print(f"\nInput: {text[:100]}...")
|
||||
|
||||
# First extract entities
|
||||
entities = self.__class__.extractor.extract_entities(text)
|
||||
self.assertTrue(len(entities) > 0, "Should extract entities first")
|
||||
|
||||
# Extract relations
|
||||
start_time = time.time()
|
||||
relations = self.__class__.relation_extractor.extract_relations(text, entities)
|
||||
duration = time.time() - start_time
|
||||
|
||||
print(f"Extracted {len(relations)} relations in {duration:.4f}s")
|
||||
for r in relations:
|
||||
print(f" - {r.subject.text} -> {r.predicate} -> {r.object.text}")
|
||||
|
||||
self.assertTrue(len(relations) > 0, "Should extract relations")
|
||||
|
||||
# Verify specific relation (Microsoft -> acquired -> Activision Blizzard)
|
||||
found_acquisition = False
|
||||
for r in relations:
|
||||
if "Microsoft" in r.subject.text and "Activision" in r.object.text:
|
||||
found_acquisition = True
|
||||
break
|
||||
|
||||
if not found_acquisition:
|
||||
# Fallback check - sometimes subject/object might be swapped or different wording
|
||||
for r in relations:
|
||||
if "Activision" in r.subject.text and "Microsoft" in r.object.text:
|
||||
found_acquisition = True
|
||||
break
|
||||
|
||||
self.assertTrue(found_acquisition, "Should find acquisition relation between Microsoft and Activision")
|
||||
|
||||
def test_5_triplet_extraction(self):
|
||||
"""Test RDF Triplet Extraction capabilities"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 5: Triplet Extraction")
|
||||
print("="*60)
|
||||
|
||||
text = self.sample_texts[0] # Apple text
|
||||
print(f"\nInput: {text[:100]}...")
|
||||
|
||||
# Pipeline: Entities -> Relations -> Triplets
|
||||
entities = self.__class__.extractor.extract_entities(text)
|
||||
relations = self.__class__.relation_extractor.extract_relations(text, entities)
|
||||
|
||||
start_time = time.time()
|
||||
triplets = self.__class__.triplet_extractor.extract_triplets(text, entities, relations)
|
||||
duration = time.time() - start_time
|
||||
|
||||
print(f"Extracted {len(triplets)} triplets in {duration:.4f}s")
|
||||
for t in triplets:
|
||||
print(f" - <{t.subject}> <{t.predicate}> <{t.object}>")
|
||||
|
||||
self.assertTrue(len(triplets) > 0, "Should extract triplets")
|
||||
|
||||
# Check for Apple related triplet
|
||||
found_apple = False
|
||||
for t in triplets:
|
||||
if "Apple" in t.subject or "Apple" in t.object:
|
||||
found_apple = True
|
||||
break
|
||||
self.assertTrue(found_apple, "Should find Apple-related triplet")
|
||||
|
||||
def test_6_event_detection(self):
|
||||
"""Test Event Detection capabilities"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 6: Event Detection")
|
||||
print("="*60)
|
||||
|
||||
text = self.sample_texts[2] # SpaceX launch text
|
||||
print(f"\nInput: {text[:100]}...")
|
||||
|
||||
start_time = time.time()
|
||||
events = self.__class__.event_detector.detect_events(text)
|
||||
duration = time.time() - start_time
|
||||
|
||||
print(f"Detected {len(events)} events in {duration:.4f}s")
|
||||
for e in events:
|
||||
print(f" - [{e.event_type}] {e.text} (Participants: {e.participants})")
|
||||
|
||||
self.assertTrue(len(events) > 0, "Should detect events")
|
||||
|
||||
# Verify launch event
|
||||
found_launch = False
|
||||
for e in events:
|
||||
if "launch" in e.event_type.lower() or "launch" in e.text.lower():
|
||||
found_launch = True
|
||||
break
|
||||
self.assertTrue(found_launch, "Should detect launch event")
|
||||
|
||||
def test_7_semantic_network(self):
|
||||
"""Test Semantic Network Extraction capabilities"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 7: Semantic Network Extraction")
|
||||
print("="*60)
|
||||
|
||||
text = self.sample_texts[3] # Google DeepMind text
|
||||
print(f"\nInput: {text[:100]}...")
|
||||
|
||||
# Extract base components first
|
||||
entities = self.__class__.extractor.extract_entities(text)
|
||||
relations = self.__class__.relation_extractor.extract_relations(text, entities)
|
||||
|
||||
start_time = time.time()
|
||||
network = self.__class__.network_extractor.extract_network(text, entities=entities, relations=relations)
|
||||
duration = time.time() - start_time
|
||||
|
||||
print(f"Extracted Network in {duration:.4f}s")
|
||||
print(f" - Nodes: {len(network.nodes)}")
|
||||
print(f" - Edges: {len(network.edges)}")
|
||||
|
||||
self.assertTrue(len(network.nodes) > 0, "Should have nodes")
|
||||
self.assertTrue(len(network.edges) > 0, "Should have edges")
|
||||
|
||||
# Verify Google/DeepMind/Gemini nodes exist
|
||||
node_labels = [n.label for n in network.nodes]
|
||||
print(f" - Node Labels: {node_labels}")
|
||||
self.assertTrue(any("Gemini" in l for l in node_labels), "Should contain Gemini node")
|
||||
|
||||
def test_8_parallel_event_detection(self):
|
||||
"""Test Parallel Event Detection capabilities"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 8: Parallel Event Detection")
|
||||
print("="*60)
|
||||
|
||||
# Create a larger batch by duplicating sample texts
|
||||
batch_texts = self.sample_texts * 2 # 10 documents
|
||||
|
||||
# 1. Sequential Run
|
||||
print("\nStarting Sequential Event Detection (10 documents)...")
|
||||
start_time = time.time()
|
||||
seq_results = self.__class__.event_detector.extract(batch_texts, max_workers=1)
|
||||
seq_time = time.time() - start_time
|
||||
print(f"Sequential Time: {seq_time:.4f}s")
|
||||
|
||||
# 2. Parallel Run
|
||||
print("\nStarting Parallel Event Detection (10 documents, 5 workers)...")
|
||||
start_time = time.time()
|
||||
par_results = self.__class__.event_detector.extract(batch_texts, max_workers=5)
|
||||
par_time = time.time() - start_time
|
||||
print(f"Parallel Time: {par_time:.4f}s")
|
||||
|
||||
# Analysis
|
||||
speedup = seq_time / par_time if par_time > 0 else 0
|
||||
print(f"\n>>> Performance Gain: {speedup:.2f}x speedup")
|
||||
|
||||
self.assertEqual(len(seq_results), len(batch_texts))
|
||||
self.assertEqual(len(par_results), len(batch_texts))
|
||||
|
||||
# Verify results match (order should be preserved)
|
||||
for i in range(len(batch_texts)):
|
||||
self.assertEqual(len(seq_results[i]), len(par_results[i]), f"Result count mismatch at index {i}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,63 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
pytest.importorskip("groq")
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
|
||||
|
||||
def test_groq_llm_smoke_entities_relations_triplets():
|
||||
if not os.getenv("GROQ_API_KEY"):
|
||||
pytest.skip("GROQ_API_KEY is not set")
|
||||
|
||||
text = (
|
||||
"Apple acquired Beats in 2014 for $3 billion. "
|
||||
"Steve Jobs founded Apple. "
|
||||
"Beats is based in California."
|
||||
)
|
||||
model = "llama-3.3-70b-versatile"
|
||||
|
||||
entities = NERExtractor(method="llm").extract(
|
||||
text,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=250,
|
||||
)
|
||||
assert isinstance(entities, list)
|
||||
assert len(entities) > 0
|
||||
assert len(entities) <= 30
|
||||
|
||||
relations = RelationExtractor(method="llm").extract(
|
||||
text,
|
||||
entities=entities,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=350,
|
||||
max_entities_prompt=12,
|
||||
)
|
||||
assert isinstance(relations, list)
|
||||
assert len(relations) <= 30
|
||||
|
||||
triplets = TripletExtractor(method="llm").extract(
|
||||
text,
|
||||
entities=entities,
|
||||
relations=relations,
|
||||
provider="groq",
|
||||
model=model,
|
||||
temperature=0.0,
|
||||
max_tokens=350,
|
||||
)
|
||||
assert isinstance(triplets, list)
|
||||
assert len(triplets) <= 40
|
||||
@@ -0,0 +1,286 @@
|
||||
import time
|
||||
import unittest
|
||||
print("Starting tests module...")
|
||||
from unittest.mock import MagicMock, patch
|
||||
from semantica.semantic_extract.providers import create_provider, ProviderPool, _provider_pool
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor, Triplet
|
||||
from semantica.semantic_extract.methods import _result_cache, extract_entities_llm, extract_relations_llm, extract_triplets_llm, match_entity
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
|
||||
class TestSemanticExtractImprovements(unittest.TestCase):
|
||||
def setUp(self):
|
||||
_provider_pool.clear()
|
||||
# Clear cache before each test
|
||||
if _result_cache:
|
||||
_result_cache._caches["entities"].clear()
|
||||
_result_cache._caches["relations"].clear()
|
||||
_result_cache._caches["triplets"].clear()
|
||||
|
||||
def test_entity_matching(self):
|
||||
print("\nTesting Entity Matching...")
|
||||
entities = [
|
||||
Entity(text="Apple Inc.", label="ORG", start_char=0, end_char=10, confidence=1.0),
|
||||
Entity(text="Steve Jobs", label="PERSON", start_char=0, end_char=10, confidence=1.0)
|
||||
]
|
||||
|
||||
# Exact match
|
||||
m1 = match_entity("Apple Inc.", entities)
|
||||
self.assertIsNotNone(m1)
|
||||
self.assertEqual(m1.text, "Apple Inc.")
|
||||
|
||||
# Case insensitive
|
||||
m2 = match_entity("apple inc.", entities)
|
||||
self.assertIsNotNone(m2)
|
||||
self.assertEqual(m2.text, "Apple Inc.")
|
||||
|
||||
# Substring/Partial match (should work via calculate_similarity)
|
||||
# "Apple" is contained in "Apple Inc."
|
||||
# calculate_similarity gives a boost for containment
|
||||
m3 = match_entity("Apple", entities)
|
||||
if m3:
|
||||
self.assertEqual(m3.text, "Apple Inc.")
|
||||
print(" Partial match 'Apple' -> 'Apple Inc.' successful.")
|
||||
else:
|
||||
print(" Partial match 'Apple' -> 'Apple Inc.' failed (score too low).")
|
||||
|
||||
# No match
|
||||
m4 = match_entity("Microsoft", entities)
|
||||
self.assertIsNone(m4)
|
||||
print(" No match verified.")
|
||||
|
||||
# Synonym match
|
||||
# We need entities that match the synonym keys in methods.py (e.g. "acquired" -> "bought")
|
||||
# Let's create an entity "bought"
|
||||
rel_entities = [Entity(text="bought", label="RELATION", start_char=0, end_char=6, confidence=1.0)]
|
||||
m5 = match_entity("acquired", rel_entities)
|
||||
self.assertIsNotNone(m5)
|
||||
self.assertEqual(m5.text, "bought")
|
||||
print(" Synonym match 'acquired' -> 'bought' verified.")
|
||||
|
||||
# Empty input
|
||||
m6 = match_entity("", entities)
|
||||
self.assertIsNone(m6)
|
||||
print(" Empty input handled.")
|
||||
|
||||
def test_caching(self):
|
||||
print("\nTesting Caching...")
|
||||
|
||||
text = "Apple Inc. was founded in 1976."
|
||||
|
||||
# Mock provider
|
||||
mock_provider = MagicMock()
|
||||
mock_provider.is_available.return_value = True
|
||||
# Setup mock response for entities
|
||||
mock_entities_response = MagicMock()
|
||||
mock_entities_response.entities = [
|
||||
MagicMock(text="Apple Inc.", label="ORG", confidence=0.9),
|
||||
MagicMock(text="1976", label="DATE", confidence=0.9)
|
||||
]
|
||||
mock_provider.generate_typed.return_value = mock_entities_response
|
||||
|
||||
with patch('semantica.semantic_extract.methods.create_provider', return_value=mock_provider) as mock_create:
|
||||
# First call - should hit provider
|
||||
print(" First call (cache miss)...")
|
||||
results1 = extract_entities_llm(text, provider="openai", model="gpt-4", api_key="test")
|
||||
self.assertEqual(len(results1), 2)
|
||||
self.assertEqual(mock_provider.generate_typed.call_count, 1)
|
||||
|
||||
# Check cache state
|
||||
print(f" Cache size: {len(_result_cache._caches['entities'])}")
|
||||
|
||||
# Second call - should hit cache
|
||||
print(" Second call (cache hit)...")
|
||||
results2 = extract_entities_llm(text, provider="openai", model="gpt-4", api_key="test")
|
||||
self.assertEqual(len(results2), 2)
|
||||
# Provider should NOT be called again
|
||||
self.assertEqual(mock_provider.generate_typed.call_count, 1)
|
||||
|
||||
print(" Cache hit verified for entities.")
|
||||
|
||||
def test_secure_caching(self):
|
||||
"""Test that sensitive parameters are excluded from cache keys."""
|
||||
print("\nTesting Secure Caching...")
|
||||
|
||||
text = "Security test."
|
||||
|
||||
# Mock provider
|
||||
mock_provider = MagicMock()
|
||||
mock_provider.is_available.return_value = True
|
||||
mock_entities_response = MagicMock()
|
||||
mock_entities_response.entities = [MagicMock(text="Test", label="TEST", confidence=1.0)]
|
||||
mock_provider.generate_typed.return_value = mock_entities_response
|
||||
|
||||
with patch('semantica.semantic_extract.methods.create_provider', return_value=mock_provider):
|
||||
# First call with one API key
|
||||
extract_entities_llm(text, provider="openai", model="gpt-4", api_key="secret_key_1")
|
||||
|
||||
# Second call with DIFFERENT API key
|
||||
# If secure caching is working, this should be a CACHE HIT because api_key is ignored
|
||||
extract_entities_llm(text, provider="openai", model="gpt-4", api_key="secret_key_2")
|
||||
|
||||
# Provider should have been called ONLY ONCE
|
||||
self.assertEqual(mock_provider.generate_typed.call_count, 1)
|
||||
print(" Secure caching verified: Changing API key did not trigger new extraction.")
|
||||
|
||||
# Verify cache content
|
||||
self.assertIn("entities", _result_cache._caches)
|
||||
self.assertTrue(len(_result_cache._caches["entities"]) > 0)
|
||||
|
||||
def test_provider_pool(self):
|
||||
print("\nTesting Provider Pool...")
|
||||
# Create provider twice with same args
|
||||
# We need to mock the actual provider init to avoid API keys requirement if not present
|
||||
with patch('semantica.semantic_extract.providers.OpenAIProvider') as MockProvider:
|
||||
MockProvider.side_effect = lambda *args, **kwargs: MagicMock()
|
||||
|
||||
p1 = create_provider("openai", api_key="test", model_name="gpt-4")
|
||||
p2 = create_provider("openai", api_key="test", model_name="gpt-4")
|
||||
|
||||
# Should be same instance
|
||||
self.assertIs(p1, p2)
|
||||
print(" Provider reuse verified.")
|
||||
|
||||
# Different args
|
||||
p3 = create_provider("openai", api_key="test", model_name="gpt-3.5")
|
||||
self.assertIsNot(p1, p3)
|
||||
print(" Different args create new instance verified.")
|
||||
|
||||
# Explicitly not using pool
|
||||
p4 = create_provider("openai", use_pool=False, api_key="test", model_name="gpt-4")
|
||||
self.assertIsNot(p1, p4)
|
||||
print(" Opt-out of pool verified.")
|
||||
|
||||
def test_ner_parallel_processing(self):
|
||||
print("\nTesting NER Parallel Processing...")
|
||||
extractor = NERExtractor(method="pattern") # Use pattern which is fast/local
|
||||
|
||||
# Mock extract_entities to simulate work and track thread execution
|
||||
original_extract = extractor.extract_entities
|
||||
|
||||
def mock_extract(text, **kwargs):
|
||||
time.sleep(0.1) # Simulate delay
|
||||
return original_extract(text, **kwargs)
|
||||
|
||||
extractor.extract_entities = mock_extract
|
||||
|
||||
texts = ["Text 1", "Text 2", "Text 3", "Text 4"]
|
||||
|
||||
start_time = time.time()
|
||||
results = extractor.extract(texts)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel NER (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
# Verify sequential fallback
|
||||
start_time_seq = time.time()
|
||||
extractor.extract(texts, max_workers=1)
|
||||
end_time_seq = time.time()
|
||||
duration_seq = end_time_seq - start_time_seq
|
||||
print(f" Sequential NER took {duration_seq:.4f}s")
|
||||
|
||||
# Check if parallel was indeed parallel (faster)
|
||||
# With 0.1s sleep * 4 items:
|
||||
# Sequential ~ 0.4s
|
||||
# Parallel (2 workers) ~ 0.2s + overhead
|
||||
self.assertLess(duration, duration_seq * 0.8)
|
||||
print(" Parallel execution speedup verified.")
|
||||
|
||||
def test_relation_parallel_processing(self):
|
||||
print("\nTesting Relation Parallel Processing...")
|
||||
extractor = RelationExtractor(method="pattern")
|
||||
|
||||
# Mock extract_relations
|
||||
original_extract = extractor.extract_relations
|
||||
def mock_extract(text, entities, **kwargs):
|
||||
time.sleep(0.1)
|
||||
return original_extract(text, entities, **kwargs)
|
||||
extractor.extract_relations = mock_extract
|
||||
|
||||
texts = ["Text 1", "Text 2", "Text 3", "Text 4"]
|
||||
entities = [[], [], [], []]
|
||||
|
||||
start_time = time.time()
|
||||
results = extractor.extract(texts, entities)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel RE (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
# Sequential
|
||||
start_time_seq = time.time()
|
||||
extractor.extract(texts, entities, max_workers=1)
|
||||
end_time_seq = time.time()
|
||||
duration_seq = end_time_seq - start_time_seq
|
||||
print(f" Sequential RE took {duration_seq:.4f}s")
|
||||
|
||||
self.assertLess(duration, duration_seq * 0.8)
|
||||
print(" Parallel execution speedup verified.")
|
||||
|
||||
def test_relation_extraction_fuzzy_matching(self):
|
||||
print("\nTesting Relation Extraction Fuzzy Matching...")
|
||||
extractor = RelationExtractor(method="pattern")
|
||||
|
||||
# Entities have formal names
|
||||
entities = [
|
||||
Entity(text="Apple Inc.", label="ORG", start_char=0, end_char=10, confidence=1.0),
|
||||
Entity(text="Steve Jobs", label="PERSON", start_char=21, end_char=31, confidence=1.0)
|
||||
]
|
||||
|
||||
# Text uses informal name "Apple"
|
||||
text = "Apple was founded by Steve Jobs."
|
||||
|
||||
relations = extractor.extract(text, entities)
|
||||
|
||||
found = False
|
||||
for rel in relations:
|
||||
# Check if subject matches "Apple Inc." even though text said "Apple"
|
||||
if rel.subject.text == "Apple Inc." and rel.object.text == "Steve Jobs" and rel.predicate == "founded_by":
|
||||
found = True
|
||||
print(" Successfully matched 'Apple' -> 'Apple Inc.' in relation extraction.")
|
||||
break
|
||||
|
||||
self.assertTrue(found, "Failed to extract relation with fuzzy entity matching")
|
||||
|
||||
def test_triplet_parallel_processing(self):
|
||||
print("\nTesting Triplet Parallel Processing...")
|
||||
extractor = TripletExtractor(method="pattern")
|
||||
|
||||
# Mock extract_triplets
|
||||
original_extract = extractor.extract_triplets
|
||||
def mock_extract(text, **kwargs):
|
||||
time.sleep(0.1)
|
||||
# Return dummy triplets to avoid actual extraction overhead
|
||||
return [Triplet(subject="s", predicate="p", object="o")]
|
||||
extractor.extract_triplets = mock_extract
|
||||
|
||||
texts = ["Text 1", "Text 2", "Text 3", "Text 4"]
|
||||
|
||||
start_time = time.time()
|
||||
results = extractor.extract(texts)
|
||||
end_time = time.time()
|
||||
|
||||
duration = end_time - start_time
|
||||
print(f" Parallel TE (default workers) took {duration:.4f}s")
|
||||
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
# Sequential
|
||||
start_time_seq = time.time()
|
||||
extractor.extract(texts, max_workers=1)
|
||||
end_time_seq = time.time()
|
||||
duration_seq = end_time_seq - start_time_seq
|
||||
print(f" Sequential TE took {duration_seq:.4f}s")
|
||||
|
||||
self.assertLess(duration, duration_seq * 0.8)
|
||||
print(" Parallel execution speedup verified.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
suite = unittest.TestLoader().loadTestsFromTestCase(TestSemanticExtractImprovements)
|
||||
unittest.TextTestRunner(verbosity=2).run(suite)
|
||||
@@ -0,0 +1,157 @@
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Ensure we test the local code, not the installed package
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
from semantica.semantic_extract.relation_extractor import Relation
|
||||
from semantica.semantic_extract.triplet_extractor import Triplet
|
||||
from semantica.semantic_extract.methods import extract_entities_llm
|
||||
# We import providers later inside tests to allow patching
|
||||
|
||||
class TestSemanticClasses:
|
||||
"""Test that core semantic classes do not have hardcoded max lengths."""
|
||||
|
||||
def test_entity_no_max_length(self):
|
||||
long_text = "a" * 10000
|
||||
entity = Entity(text=long_text, label="TEST", start_char=0, end_char=10000)
|
||||
assert entity.text == long_text
|
||||
assert len(entity.text) == 10000
|
||||
|
||||
def test_relation_no_max_length(self):
|
||||
long_text = "a" * 10000
|
||||
e1 = Entity(text="s", label="S", start_char=0, end_char=1)
|
||||
e2 = Entity(text="o", label="O", start_char=0, end_char=1)
|
||||
relation = Relation(subject=e1, predicate=long_text, object=e2)
|
||||
assert relation.predicate == long_text
|
||||
|
||||
def test_triplet_no_max_length(self):
|
||||
long_text = "a" * 10000
|
||||
triplet = Triplet(subject=long_text, predicate="r", object="t")
|
||||
assert triplet.subject == long_text
|
||||
|
||||
|
||||
class TestProviderLimits:
|
||||
"""Test that providers pass through correct length parameters."""
|
||||
|
||||
def test_openai_max_completion_tokens(self):
|
||||
from semantica.semantic_extract.providers import OpenAIProvider
|
||||
|
||||
# Patch _init_client to avoid real client creation and import issues
|
||||
with patch.object(OpenAIProvider, '_init_client', return_value=None):
|
||||
provider = OpenAIProvider(api_key="fake")
|
||||
|
||||
# Manually mock client
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices[0].message.content = "result"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
provider.client = mock_client
|
||||
|
||||
provider.generate("prompt", max_completion_tokens=12345, top_p=0.9)
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
assert call_kwargs["max_completion_tokens"] == 12345
|
||||
assert call_kwargs["top_p"] == 0.9
|
||||
assert "max_tokens" not in call_kwargs
|
||||
|
||||
def test_anthropic_max_tokens_defaults(self):
|
||||
from semantica.semantic_extract.providers import AnthropicProvider
|
||||
|
||||
with patch.object(AnthropicProvider, '_init_client', return_value=None):
|
||||
provider = AnthropicProvider(api_key="fake")
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = [MagicMock(text="result")]
|
||||
mock_client.messages.create.return_value = mock_response
|
||||
provider.client = mock_client
|
||||
|
||||
provider.generate("prompt")
|
||||
|
||||
# Verify default is 8192 (new limit)
|
||||
call_kwargs = mock_client.messages.create.call_args[1]
|
||||
assert call_kwargs["max_tokens"] == 8192
|
||||
|
||||
# Test override
|
||||
provider.generate("prompt", max_tokens=9999)
|
||||
call_kwargs = mock_client.messages.create.call_args[1]
|
||||
assert call_kwargs["max_tokens"] == 9999
|
||||
|
||||
def test_groq_max_completion_tokens(self):
|
||||
from semantica.semantic_extract.providers import GroqProvider
|
||||
|
||||
with patch.object(GroqProvider, '_init_client', return_value=None):
|
||||
provider = GroqProvider(api_key="fake")
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices[0].message.content = "result"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
provider.client = mock_client
|
||||
|
||||
provider.generate("prompt", max_completion_tokens=5000)
|
||||
|
||||
# Verify
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
assert call_kwargs["max_completion_tokens"] == 5000
|
||||
|
||||
def test_gemini_params(self):
|
||||
from semantica.semantic_extract.providers import GeminiProvider
|
||||
|
||||
with patch.object(GeminiProvider, '_init_client', return_value=None):
|
||||
provider = GeminiProvider(api_key="fake")
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.text = "result"
|
||||
mock_model.generate_content.return_value = mock_response
|
||||
provider.client = mock_model
|
||||
|
||||
provider.generate("prompt", top_k=10, candidate_count=2)
|
||||
|
||||
# Verify
|
||||
call_kwargs = mock_model.generate_content.call_args[1]
|
||||
gen_config = call_kwargs["generation_config"]
|
||||
assert gen_config["top_k"] == 10
|
||||
assert gen_config["candidate_count"] == 2
|
||||
|
||||
class TestChunkingDefaults:
|
||||
"""Test that chunking defaults have been increased."""
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
@patch("semantica.semantic_extract.methods._extract_entities_chunked")
|
||||
def test_openai_chunking_limit(self, mock_chunked, mock_create_provider):
|
||||
# Setup
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.is_available.return_value = True
|
||||
mock_create_provider.return_value = mock_llm
|
||||
|
||||
# Text length = 10000 (Greater than old 4000, less than new 64000)
|
||||
long_text = "a" * 10000
|
||||
|
||||
# Call without explicit max_text_length
|
||||
extract_entities_llm(long_text, provider="openai", api_key="fake")
|
||||
|
||||
# Should NOT call chunked extraction because default is now 64000
|
||||
mock_chunked.assert_not_called()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
@patch("semantica.semantic_extract.methods._extract_entities_chunked")
|
||||
def test_groq_chunking_limit(self, mock_chunked, mock_create_provider):
|
||||
# Setup
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.is_available.return_value = True
|
||||
mock_create_provider.return_value = mock_llm
|
||||
|
||||
# Text length = 10000 (Greater than old 8000, less than new 64000)
|
||||
long_text = "a" * 10000
|
||||
|
||||
extract_entities_llm(long_text, provider="groq", api_key="fake")
|
||||
|
||||
# Should NOT call chunked extraction because default is now 64000
|
||||
mock_chunked.assert_not_called()
|
||||
@@ -0,0 +1,123 @@
|
||||
import multiprocessing
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantica.semantic_extract.config import resolve_max_workers
|
||||
from semantica.semantic_extract.ner_extractor import Entity, NERExtractor
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
|
||||
from semantica.semantic_extract.methods import filter_entities_for_text
|
||||
from semantica.semantic_extract.schemas import RelationsResponse, RelationOut
|
||||
|
||||
|
||||
def test_resolve_max_workers_defaults_and_clamps():
|
||||
cpu_count = multiprocessing.cpu_count() or 1
|
||||
|
||||
assert resolve_max_workers(explicit=0) == 1
|
||||
assert resolve_max_workers(explicit=-10) == 1
|
||||
assert resolve_max_workers(explicit=1) == 1
|
||||
assert resolve_max_workers(explicit=10**9) == min(cpu_count, 32)
|
||||
|
||||
assert resolve_max_workers(explicit=None, methods=["ml"]) == 1
|
||||
|
||||
|
||||
def test_filter_entities_for_text_keeps_short_tokens():
|
||||
text = "US AI lab in NY"
|
||||
entities = [
|
||||
Entity(text="US", label="GPE", start_char=0, end_char=2, confidence=1.0),
|
||||
Entity(text="AI", label="TECH", start_char=3, end_char=5, confidence=1.0),
|
||||
Entity(text="NY", label="GPE", start_char=13, end_char=15, confidence=1.0),
|
||||
]
|
||||
kept = filter_entities_for_text(text, entities, max_keep=2)
|
||||
kept_texts = {e.text for e in kept}
|
||||
assert "US" in kept_texts or "AI" in kept_texts or "NY" in kept_texts
|
||||
|
||||
|
||||
def test_pattern_batch_defaults_to_single_worker_low_latency():
|
||||
extractor = NERExtractor(method="pattern")
|
||||
texts = [f"Text {i}" for i in range(8)]
|
||||
extractor.extract(texts)
|
||||
|
||||
|
||||
def test_relation_llm_prompt_filter_does_not_break_mapping():
|
||||
entities = [Entity(text=f"VeryLongEntityName{i}", label="ORG", start_char=0, end_char=1, confidence=1.0) for i in range(120)]
|
||||
ghost = Entity(text="Ghost", label="ORG", start_char=0, end_char=1, confidence=1.0)
|
||||
entities.append(ghost)
|
||||
|
||||
captured = {}
|
||||
|
||||
class FakeLLM:
|
||||
def is_available(self):
|
||||
return True
|
||||
|
||||
def generate_typed(self, prompt, schema, **kwargs):
|
||||
captured["prompt"] = prompt
|
||||
return RelationsResponse(
|
||||
relations=[
|
||||
RelationOut(subject="Ghost", predicate="related_to", object="VeryLongEntityName0", confidence=0.9)
|
||||
]
|
||||
)
|
||||
|
||||
with patch("semantica.semantic_extract.methods.create_provider", return_value=FakeLLM()):
|
||||
from semantica.semantic_extract.methods import extract_relations_llm
|
||||
|
||||
relations = extract_relations_llm(
|
||||
"Short text mentioning VeryLongEntityName0 only.",
|
||||
entities=entities,
|
||||
provider="openai",
|
||||
model="gpt-4",
|
||||
max_entities_prompt=20,
|
||||
)
|
||||
|
||||
assert "Ghost" not in captured["prompt"]
|
||||
assert len(relations) == 1
|
||||
assert relations[0].subject.text == "Ghost"
|
||||
|
||||
|
||||
def test_triplet_extractor_reuses_sub_extractors():
|
||||
ner_instance = MagicMock()
|
||||
ner_instance.extract_entities.return_value = [
|
||||
Entity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)
|
||||
]
|
||||
|
||||
rel_instance = MagicMock()
|
||||
rel_instance.extract_relations.return_value = []
|
||||
|
||||
ner_ctor = MagicMock(return_value=ner_instance)
|
||||
rel_ctor = MagicMock(return_value=rel_instance)
|
||||
|
||||
with patch("semantica.semantic_extract.ner_extractor.NERExtractor", ner_ctor), patch(
|
||||
"semantica.semantic_extract.relation_extractor.RelationExtractor", rel_ctor
|
||||
), patch("semantica.semantic_extract.methods.get_triplet_method", return_value=lambda *args, **kwargs: []):
|
||||
extractor = TripletExtractor(method="pattern")
|
||||
extractor.extract_triplets("A text.")
|
||||
extractor.extract_triplets("A text again.")
|
||||
|
||||
assert ner_ctor.call_count == 1
|
||||
assert rel_ctor.call_count == 1
|
||||
|
||||
|
||||
def test_semantic_network_extractor_reuses_sub_extractors():
|
||||
ner_instance = MagicMock()
|
||||
ner_instance.extract_entities.return_value = [
|
||||
Entity(text="A", label="PERSON", start_char=0, end_char=1, confidence=1.0)
|
||||
]
|
||||
|
||||
rel_instance = MagicMock()
|
||||
rel_instance.extract_relations.return_value = []
|
||||
|
||||
ner_ctor = MagicMock(return_value=ner_instance)
|
||||
rel_ctor = MagicMock(return_value=rel_instance)
|
||||
|
||||
with patch("semantica.semantic_extract.ner_extractor.NERExtractor", ner_ctor), patch(
|
||||
"semantica.semantic_extract.relation_extractor.RelationExtractor", rel_ctor
|
||||
):
|
||||
extractor = SemanticNetworkExtractor(method="pattern")
|
||||
extractor.extract_network("A text.")
|
||||
extractor.extract_network("A text again.")
|
||||
|
||||
assert ner_ctor.call_count == 1
|
||||
assert rel_ctor.call_count == 1
|
||||
@@ -0,0 +1,151 @@
|
||||
|
||||
import pytest
|
||||
import sys
|
||||
|
||||
from semantica.semantic_extract.ner_extractor import NERExtractor, Entity
|
||||
from semantica.semantic_extract.relation_extractor import RelationExtractor, Relation
|
||||
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||
|
||||
class TestRobustnessFallback:
|
||||
|
||||
def test_ner_last_resort_fallback(self):
|
||||
"""Test that NER extractor finds entities even in obscure text via last resort."""
|
||||
extractor = NERExtractor()
|
||||
|
||||
# Text with single capitalized word that shouldn't match PERSON pattern (requires 2+ words)
|
||||
text = "Zylophone"
|
||||
|
||||
entities = extractor.extract_entities(text)
|
||||
|
||||
assert len(entities) > 0, "Should have extracted at least one entity via last resort"
|
||||
# Check if they are the capitalized words
|
||||
texts = [e.text for e in entities]
|
||||
assert "Zylophone" in texts
|
||||
|
||||
# Verify metadata
|
||||
for e in entities:
|
||||
assert e.metadata is not None
|
||||
assert "extraction_method" in e.metadata
|
||||
# Should be last_resort_pattern
|
||||
assert e.metadata["extraction_method"] == "last_resort_pattern"
|
||||
|
||||
def test_relation_last_resort_fallback(self):
|
||||
"""Test that Relation extractor creates adjacency relations when no patterns match."""
|
||||
extractor = RelationExtractor()
|
||||
|
||||
# Create entities far apart to avoid "co_occurrence" fallback which triggers < 100 chars
|
||||
padding = " " * 105
|
||||
text = f"Alpha{padding}Beta{padding}Gamma"
|
||||
|
||||
# Alpha at start
|
||||
e1_start = 0
|
||||
e1_end = 5
|
||||
|
||||
# Beta after padding
|
||||
e2_start = e1_end + 105
|
||||
e2_end = e2_start + 4
|
||||
|
||||
# Gamma after padding
|
||||
e3_start = e2_end + 105
|
||||
e3_end = e3_start + 5
|
||||
|
||||
e1 = Entity(text="Alpha", label="UNKNOWN", start_char=e1_start, end_char=e1_end)
|
||||
e2 = Entity(text="Beta", label="UNKNOWN", start_char=e2_start, end_char=e2_end)
|
||||
e3 = Entity(text="Gamma", label="UNKNOWN", start_char=e3_start, end_char=e3_end)
|
||||
|
||||
entities = [e1, e2, e3]
|
||||
|
||||
# This text has no "is a", "works for", etc. patterns.
|
||||
# And entities are too far for co-occurrence (< 100).
|
||||
# It should trigger the last resort adjacency fallback.
|
||||
relations = extractor.extract_relations(text, entities)
|
||||
|
||||
assert len(relations) > 0, "Should have extracted relations via last resort"
|
||||
|
||||
# Expect relations between adjacent entities: Alpha->Beta, Beta->Gamma
|
||||
pairs = [(r.subject.text, r.object.text) for r in relations]
|
||||
assert ("Alpha", "Beta") in pairs
|
||||
assert ("Beta", "Gamma") in pairs
|
||||
|
||||
# Verify metadata
|
||||
for r in relations:
|
||||
assert r.metadata is not None
|
||||
assert "extraction_method" in r.metadata
|
||||
assert r.metadata.get("extraction_method") == "last_resort_adjacency"
|
||||
|
||||
def test_triplet_fallback_conversion(self):
|
||||
"""Test that Triplet extractor falls back to converting relations if extraction fails."""
|
||||
# Setup mocks or use real classes
|
||||
ner = NERExtractor() # We'll just pass entities directly
|
||||
rel_extractor = RelationExtractor()
|
||||
triplet_extractor = TripletExtractor()
|
||||
|
||||
text = "Alpha is connected to Beta."
|
||||
e1 = Entity(text="Alpha", label="Thing", start_char=0, end_char=5)
|
||||
e2 = Entity(text="Beta", label="Thing", start_char=22, end_char=26)
|
||||
entities = [e1, e2]
|
||||
|
||||
# Create a relation manually to ensure we have one to convert
|
||||
relation = Relation(
|
||||
subject=e1,
|
||||
predicate="connected_to",
|
||||
object=e2,
|
||||
confidence=0.9,
|
||||
context=text
|
||||
)
|
||||
|
||||
# We want to test the fallback in extract_triplets.
|
||||
# Since we can't easily force the primary triplet method to return empty without mocking,
|
||||
# we can pass the relations explicitly and rely on the fact that standard triplet extraction
|
||||
# might not support "connected_to" if it relies on strict patterns, or we can use a method that fails.
|
||||
|
||||
# However, the triplet extractor calls relation extractor internally if not provided.
|
||||
# Let's test the flow where we provide relations.
|
||||
|
||||
triplets = triplet_extractor.extract_triplets(text, entities=entities, relations=[relation])
|
||||
|
||||
assert len(triplets) > 0
|
||||
assert triplets[0].subject == "Alpha"
|
||||
assert triplets[0].object == "Beta"
|
||||
assert triplets[0].predicate == "connected_to"
|
||||
|
||||
def test_batch_metadata_propagation(self):
|
||||
"""Verify batch_index and document_id are propagated in batch mode with fallbacks."""
|
||||
ner = NERExtractor()
|
||||
|
||||
docs = [
|
||||
{"content": "First doc", "id": "doc_1"},
|
||||
{"content": "Second doc", "id": "doc_2"}
|
||||
]
|
||||
|
||||
# These docs are simple, might trigger fallback or simple patterns
|
||||
results = ner.extract(docs)
|
||||
|
||||
assert len(results) == 2
|
||||
|
||||
# Check first doc results
|
||||
for e in results[0]:
|
||||
assert e.metadata["batch_index"] == 0
|
||||
assert e.metadata["document_id"] == "doc_1"
|
||||
|
||||
# Check second doc results
|
||||
for e in results[1]:
|
||||
assert e.metadata["batch_index"] == 1
|
||||
assert e.metadata["document_id"] == "doc_2"
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Manually run if executed as script
|
||||
t = TestRobustnessFallback()
|
||||
try:
|
||||
t.test_ner_last_resort_fallback()
|
||||
print("NER Fallback Test Passed")
|
||||
t.test_relation_last_resort_fallback()
|
||||
print("Relation Fallback Test Passed")
|
||||
t.test_triplet_fallback_conversion()
|
||||
print("Triplet Fallback Test Passed")
|
||||
t.test_batch_metadata_propagation()
|
||||
print("Batch Metadata Test Passed")
|
||||
except Exception as e:
|
||||
print(f"Test Failed: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
@@ -0,0 +1,177 @@
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from typing import Type, List, Optional
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from semantica.semantic_extract.providers import BaseProvider
|
||||
from semantica.semantic_extract.methods import (
|
||||
extract_entities_llm,
|
||||
extract_relations_llm,
|
||||
extract_triplets_llm
|
||||
)
|
||||
from semantica.semantic_extract.schemas import EntitiesResponse, RelationsResponse, TripletsResponse
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
from semantica.semantic_extract.relation_extractor import Relation
|
||||
|
||||
# Mock Pydantic models for responses
|
||||
class MockEntity(BaseModel):
|
||||
text: str
|
||||
label: str
|
||||
start: int = 0
|
||||
end: int = 0
|
||||
confidence: float = 1.0
|
||||
|
||||
class MockEntitiesResponse(BaseModel):
|
||||
entities: List[MockEntity]
|
||||
|
||||
class MockRelation(BaseModel):
|
||||
subject: str
|
||||
predicate: str
|
||||
object: str
|
||||
confidence: float = 1.0
|
||||
|
||||
class MockRelationsResponse(BaseModel):
|
||||
relations: List[MockRelation]
|
||||
|
||||
class MockTriplet(BaseModel):
|
||||
subject: str
|
||||
predicate: str
|
||||
object: str
|
||||
confidence: float = 1.0
|
||||
|
||||
class MockTripletsResponse(BaseModel):
|
||||
triplets: List[MockTriplet]
|
||||
|
||||
class MockProvider(BaseProvider):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.generate_typed_called = False
|
||||
self.generate_structured_called = False
|
||||
self.model = "mock-model"
|
||||
self.is_available_val = True
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.is_available_val
|
||||
|
||||
def generate(self, prompt: str, **kwargs) -> str:
|
||||
return "{}"
|
||||
|
||||
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||
self.generate_structured_called = True
|
||||
if "entities" in prompt.lower():
|
||||
return [{"text": "Apple", "label": "ORG", "start": 0, "end": 5}]
|
||||
elif "relations" in prompt.lower():
|
||||
return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}]
|
||||
elif "triplets" in prompt.lower():
|
||||
return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}]
|
||||
return {}
|
||||
|
||||
def generate_typed(
|
||||
self,
|
||||
prompt: str,
|
||||
schema: Type[BaseModel],
|
||||
max_retries: int = 3,
|
||||
**kwargs
|
||||
) -> BaseModel:
|
||||
self.generate_typed_called = True
|
||||
|
||||
if schema.__name__ == "EntitiesResponse":
|
||||
return EntitiesResponse(entities=[
|
||||
{"text": "Apple", "label": "ORG", "start_char": 0, "end_char": 5, "confidence": 0.99}
|
||||
])
|
||||
elif schema.__name__ == "RelationsResponse":
|
||||
return RelationsResponse(relations=[
|
||||
{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95}
|
||||
])
|
||||
elif schema.__name__ == "TripletsResponse":
|
||||
return TripletsResponse(triplets=[
|
||||
{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95}
|
||||
])
|
||||
return schema()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_provider():
|
||||
return MockProvider()
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_extract_entities_typed(mock_create_provider, mock_provider):
|
||||
mock_create_provider.return_value = mock_provider
|
||||
|
||||
text = "Apple was founded by Steve Jobs."
|
||||
entities = extract_entities_llm(
|
||||
text,
|
||||
provider="mock",
|
||||
structured_output_mode="typed"
|
||||
)
|
||||
|
||||
assert mock_provider.generate_typed_called
|
||||
assert len(entities) == 1
|
||||
assert entities[0].text == "Apple"
|
||||
assert entities[0].label == "ORG"
|
||||
assert entities[0].metadata["extraction_method"] == "llm_typed"
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_extract_entities_legacy(mock_create_provider, mock_provider):
|
||||
mock_create_provider.return_value = mock_provider
|
||||
|
||||
text = "Apple was founded by Steve Jobs."
|
||||
entities = extract_entities_llm(
|
||||
text,
|
||||
provider="mock",
|
||||
structured_output_mode="legacy"
|
||||
)
|
||||
|
||||
# Legacy mode now redirects to typed mode
|
||||
assert mock_provider.generate_typed_called
|
||||
assert not mock_provider.generate_structured_called
|
||||
assert len(entities) == 1
|
||||
assert entities[0].text == "Apple"
|
||||
assert entities[0].label == "ORG"
|
||||
assert entities[0].metadata["extraction_method"] == "llm_typed"
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_extract_relations_typed(mock_create_provider, mock_provider):
|
||||
mock_create_provider.return_value = mock_provider
|
||||
|
||||
text = "Steve Jobs founded Apple."
|
||||
entities = [
|
||||
Entity(text="Steve Jobs", label="PERSON", start_char=0, end_char=10),
|
||||
Entity(text="Apple", label="ORG", start_char=19, end_char=24)
|
||||
]
|
||||
|
||||
relations = extract_relations_llm(
|
||||
text,
|
||||
entities=entities,
|
||||
provider="mock",
|
||||
structured_output_mode="typed"
|
||||
)
|
||||
|
||||
assert mock_provider.generate_typed_called
|
||||
assert len(relations) == 1
|
||||
assert relations[0].subject.text == "Steve Jobs"
|
||||
assert relations[0].object.text == "Apple"
|
||||
assert relations[0].predicate == "founded"
|
||||
assert relations[0].metadata["extraction_method"] == "llm_typed"
|
||||
|
||||
@patch("semantica.semantic_extract.methods.create_provider")
|
||||
def test_extract_triplets_typed(mock_create_provider, mock_provider):
|
||||
mock_create_provider.return_value = mock_provider
|
||||
|
||||
text = "Steve Jobs founded Apple."
|
||||
|
||||
triplets = extract_triplets_llm(
|
||||
text,
|
||||
provider="mock",
|
||||
structured_output_mode="typed"
|
||||
)
|
||||
|
||||
assert mock_provider.generate_typed_called
|
||||
assert len(triplets) == 1
|
||||
assert triplets[0].subject == "Steve Jobs"
|
||||
assert triplets[0].object == "Apple"
|
||||
assert triplets[0].predicate == "founded"
|
||||
assert triplets[0].metadata["extraction_method"] == "llm_typed"
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
File diff suppressed because it is too large
Load Diff
@@ -34,7 +34,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
||||
"""Test that methods raise ProcessingError by default on failure."""
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.is_available.return_value = True
|
||||
mock_llm.generate_structured.side_effect = ProcessingError("LLM Error")
|
||||
mock_llm.generate_typed.side_effect = ProcessingError("LLM Error")
|
||||
mock_create.return_value = mock_llm
|
||||
|
||||
try:
|
||||
@@ -48,7 +48,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
||||
"""Test that silent_fail=True returns empty list instead of raising."""
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.is_available.return_value = True
|
||||
mock_llm.generate_structured.side_effect = Exception("LLM Error")
|
||||
mock_llm.generate_typed.side_effect = Exception("LLM Error")
|
||||
mock_create.return_value = mock_llm
|
||||
|
||||
entities = extract_entities_llm("test text", provider="openai", silent_fail=True)
|
||||
@@ -83,7 +83,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
||||
"""Test that long text triggers chunking."""
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.is_available.return_value = True
|
||||
mock_llm.generate_structured.return_value = []
|
||||
mock_llm.generate_typed.return_value = MagicMock(entities=[]) # Mock response
|
||||
mock_create.return_value = mock_llm
|
||||
|
||||
long_text = "This is a long text that should be chunked into multiple pieces."
|
||||
|
||||
@@ -32,7 +32,8 @@ class TestModelSelection(unittest.TestCase):
|
||||
|
||||
def test_generator_switching(self):
|
||||
print("\nTesting EmbeddingGenerator Switching...")
|
||||
generator = EmbeddingGenerator()
|
||||
# Initialize with explicit method to ensure consistent starting state for test
|
||||
generator = EmbeddingGenerator(text={"method": "sentence_transformers"})
|
||||
|
||||
# Default check
|
||||
self.assertEqual(generator.get_text_method(), "sentence_transformers")
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
import unittest
|
||||
|
||||
from semantica.semantic_extract.methods import extract_relations_llm
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
|
||||
|
||||
class FakeProvider:
|
||||
def __init__(self, typed_payload=None, structured_payload=None):
|
||||
self._typed_payload = typed_payload
|
||||
self._structured_payload = structured_payload
|
||||
|
||||
def is_available(self):
|
||||
return True
|
||||
|
||||
# Simulate typed output return: can be dict or an object with relations
|
||||
def generate_typed(self, prompt, schema, **kwargs):
|
||||
return self._typed_payload if self._typed_payload is not None else {"relations": []}
|
||||
|
||||
def generate_structured(self, prompt, **kwargs):
|
||||
return self._structured_payload if self._structured_payload is not None else {"relations": []}
|
||||
|
||||
|
||||
class TestLLMRelationExtraction(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Minimal realistic text and entities
|
||||
self.text = "Apple reported revenue of $4.4 billion in Q1 2024."
|
||||
self.entities = [
|
||||
Entity(text="Apple", label="ORGANIZATION", start_char=0, end_char=5, confidence=0.99),
|
||||
Entity(text="$4.4 billion", label="MONEY", start_char=26, end_char=39, confidence=0.99),
|
||||
Entity(text="Q1 2024", label="DATE", start_char=43, end_char=51, confidence=0.99),
|
||||
]
|
||||
|
||||
def _monkeypatch_provider(self, provider_instance):
|
||||
# Monkeypatch create_provider used by extract_relations_llm
|
||||
import semantica.semantic_extract.methods as methods
|
||||
self._orig_create_provider = methods.create_provider
|
||||
|
||||
def _fake_create_provider(provider, model=None, **kwargs):
|
||||
return provider_instance
|
||||
|
||||
methods.create_provider = _fake_create_provider
|
||||
|
||||
def tearDown(self):
|
||||
# Restore original create_provider if patched
|
||||
try:
|
||||
import semantica.semantic_extract.methods as methods
|
||||
if hasattr(self, "_orig_create_provider"):
|
||||
methods.create_provider = self._orig_create_provider
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def test_typed_relations_parsed(self):
|
||||
# Typed returns a dict compatible with parser
|
||||
typed_payload = {
|
||||
"relations": [
|
||||
{
|
||||
"subject": "Apple",
|
||||
"predicate": "HAS_REVENUE",
|
||||
"object": "$4.4 billion",
|
||||
"confidence": 0.92,
|
||||
},
|
||||
{
|
||||
"subject": "Apple",
|
||||
"predicate": "IN_QUARTER",
|
||||
"object": "Q1 2024",
|
||||
"confidence": 0.9,
|
||||
},
|
||||
]
|
||||
}
|
||||
fake = FakeProvider(typed_payload=typed_payload)
|
||||
self._monkeypatch_provider(fake)
|
||||
|
||||
rels = extract_relations_llm(
|
||||
text=self.text,
|
||||
entities=self.entities,
|
||||
provider="groq",
|
||||
model="llama-3.1-8b-instant",
|
||||
relation_types=["HAS_REVENUE", "IN_QUARTER"],
|
||||
verbose=True,
|
||||
)
|
||||
self.assertGreaterEqual(len(rels), 2, "Expected at least two relations from typed payload")
|
||||
preds = {(r.subject.text, r.predicate, r.object.text) for r in rels}
|
||||
self.assertIn(("Apple", "HAS_REVENUE", "$4.4 billion"), preds)
|
||||
self.assertIn(("Apple", "IN_QUARTER", "Q1 2024"), preds)
|
||||
|
||||
def test_structured_fallback_used(self):
|
||||
# Typed returns zero, structured has content
|
||||
typed_payload = {"relations": []}
|
||||
structured_payload = {
|
||||
"relations": [
|
||||
{
|
||||
"subject": "Apple",
|
||||
"predicate": "HAS_REVENUE",
|
||||
"object": "$4.4 billion",
|
||||
"confidence": 0.88,
|
||||
}
|
||||
]
|
||||
}
|
||||
fake = FakeProvider(typed_payload=typed_payload, structured_payload=structured_payload)
|
||||
self._monkeypatch_provider(fake)
|
||||
|
||||
rels = extract_relations_llm(
|
||||
text=self.text,
|
||||
entities=self.entities,
|
||||
provider="groq",
|
||||
model="llama-3.1-8b-instant",
|
||||
relation_types=["HAS_REVENUE"],
|
||||
verbose=True,
|
||||
)
|
||||
self.assertEqual(len(rels), 1, "Expected fallback to structured JSON to yield one relation")
|
||||
r = rels[0]
|
||||
self.assertEqual(r.subject.text, "Apple")
|
||||
self.assertEqual(r.object.text, "$4.4 billion")
|
||||
self.assertEqual(r.predicate, "HAS_REVENUE")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,122 @@
|
||||
|
||||
print("Starting tests module...")
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
try:
|
||||
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor, SemanticNetwork, SemanticNode, SemanticEdge
|
||||
from semantica.semantic_extract.event_detector import EventDetector, Event
|
||||
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||
from semantica.semantic_extract.coreference_resolver import CoreferenceResolver, CoreferenceChain, Mention
|
||||
from semantica.semantic_extract.ner_extractor import Entity
|
||||
print("Imports successful")
|
||||
except Exception as e:
|
||||
print(f"Import failed: {e}")
|
||||
|
||||
class TestSemanticExtractBatch(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Mock progress tracker to avoid console spam
|
||||
self.tracker_patcher = patch('semantica.utils.progress_tracker.get_progress_tracker')
|
||||
self.mock_tracker_cls = self.tracker_patcher.start()
|
||||
self.mock_tracker = self.mock_tracker_cls.return_value
|
||||
self.mock_tracker.enabled = True
|
||||
self.mock_tracker.start_tracking.return_value = "tracking_id"
|
||||
|
||||
def tearDown(self):
|
||||
self.tracker_patcher.stop()
|
||||
|
||||
def test_semantic_network_batch(self):
|
||||
print("Running test_semantic_network_batch")
|
||||
from copy import deepcopy
|
||||
extractor = SemanticNetworkExtractor()
|
||||
|
||||
# Mock extract_network
|
||||
mock_network = SemanticNetwork(
|
||||
nodes=[SemanticNode(id="1", label="test", type="test", metadata={})],
|
||||
edges=[SemanticEdge(source="1", target="1", label="self", metadata={})],
|
||||
metadata={}
|
||||
)
|
||||
# Use side_effect to return a fresh copy each time
|
||||
extractor.extract_network = MagicMock(side_effect=lambda *args, **kwargs: deepcopy(mock_network))
|
||||
|
||||
# Test input
|
||||
docs = [{"content": "doc1", "id": "doc_1"}, {"content": "doc2", "id": "doc_2"}]
|
||||
|
||||
# Run batch
|
||||
results = extractor.extract(docs)
|
||||
|
||||
self.assertEqual(len(results), 2)
|
||||
# Check provenance
|
||||
self.assertEqual(results[0].metadata["batch_index"], 0)
|
||||
self.assertEqual(results[0].metadata["document_id"], "doc_1")
|
||||
self.assertEqual(results[0].nodes[0].metadata["batch_index"], 0)
|
||||
self.assertEqual(results[0].nodes[0].metadata["document_id"], "doc_1")
|
||||
|
||||
self.assertEqual(results[1].metadata["batch_index"], 1)
|
||||
self.assertEqual(results[1].metadata["document_id"], "doc_2")
|
||||
|
||||
def test_event_detector_batch(self):
|
||||
print("Running test_event_detector_batch")
|
||||
detector = EventDetector()
|
||||
|
||||
# Mock detect_events
|
||||
mock_event = Event(
|
||||
text="event", event_type="test", start_char=0, end_char=5
|
||||
)
|
||||
detector.detect_events = MagicMock(return_value=[mock_event])
|
||||
|
||||
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||
|
||||
results = detector.extract(docs)
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(len(results[0]), 1)
|
||||
self.assertEqual(results[0][0].metadata["batch_index"], 0)
|
||||
self.assertEqual(results[0][0].metadata["document_id"], "doc_1")
|
||||
|
||||
def test_semantic_analyzer_batch(self):
|
||||
print("Running test_semantic_analyzer_batch")
|
||||
analyzer = SemanticAnalyzer()
|
||||
|
||||
# Mock analyze_semantics
|
||||
mock_result = {
|
||||
"text": "test",
|
||||
"semantic_roles": [{"word": "test", "role": "agent"}]
|
||||
}
|
||||
analyzer.analyze_semantics = MagicMock(return_value=mock_result)
|
||||
|
||||
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||
|
||||
results = analyzer.analyze(docs)
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["batch_index"], 0)
|
||||
self.assertEqual(results[0]["document_id"], "doc_1")
|
||||
self.assertEqual(results[0]["semantic_roles"][0]["metadata"]["batch_index"], 0)
|
||||
self.assertEqual(results[0]["semantic_roles"][0]["metadata"]["document_id"], "doc_1")
|
||||
|
||||
def test_coreference_resolver_batch(self):
|
||||
print("Running test_coreference_resolver_batch")
|
||||
resolver = CoreferenceResolver()
|
||||
|
||||
# Mock resolve_coreferences
|
||||
mock_mention = Mention(text="he", start_char=0, end_char=2, mention_type="pronoun")
|
||||
mock_chain = CoreferenceChain(
|
||||
mentions=[mock_mention],
|
||||
representative=mock_mention
|
||||
)
|
||||
resolver.resolve_coreferences = MagicMock(return_value=[mock_chain])
|
||||
|
||||
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||
|
||||
results = resolver.resolve(docs)
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0][0].mentions[0].metadata["batch_index"], 0)
|
||||
self.assertEqual(results[0][0].mentions[0].metadata["document_id"], "doc_1")
|
||||
self.assertEqual(results[0][0].representative.metadata["batch_index"], 0)
|
||||
self.assertEqual(results[0][0].representative.metadata["document_id"], "doc_1")
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("Running main...")
|
||||
unittest.main()
|
||||
@@ -1,7 +1,7 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from semantica.triplet_store.triplet_manager import TripletManager, TripletStore
|
||||
from semantica.triplet_store.query_engine import QueryEngine, QueryResult
|
||||
from semantica.triplet_store.triplet_store import TripletStore
|
||||
from semantica.triplet_store.query_engine import QueryEngine
|
||||
from semantica.semantic_extract.triplet_extractor import Triplet
|
||||
|
||||
class TestTripletStore(unittest.TestCase):
|
||||
@@ -10,125 +10,91 @@ class TestTripletStore(unittest.TestCase):
|
||||
self.mock_logger = MagicMock()
|
||||
self.mock_tracker = MagicMock()
|
||||
|
||||
self.logger_patcher = patch('semantica.triplet_store.triplet_manager.get_logger', return_value=self.mock_logger)
|
||||
self.tracker_patcher = patch('semantica.triplet_store.triplet_manager.get_progress_tracker', return_value=self.mock_tracker)
|
||||
self.logger_patcher_qe = patch('semantica.triplet_store.query_engine.get_logger', return_value=self.mock_logger)
|
||||
self.tracker_patcher_qe = patch('semantica.triplet_store.query_engine.get_progress_tracker', return_value=self.mock_tracker)
|
||||
self.logger_patcher = patch('semantica.triplet_store.triplet_store.get_logger', return_value=self.mock_logger)
|
||||
self.tracker_patcher = patch('semantica.triplet_store.triplet_store.get_progress_tracker', return_value=self.mock_tracker)
|
||||
|
||||
self.logger_patcher.start()
|
||||
self.tracker_patcher.start()
|
||||
self.logger_patcher_qe.start()
|
||||
self.tracker_patcher_qe.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.logger_patcher.stop()
|
||||
self.tracker_patcher.stop()
|
||||
self.logger_patcher_qe.stop()
|
||||
self.tracker_patcher_qe.stop()
|
||||
|
||||
def test_triplet_manager_init(self):
|
||||
manager = TripletManager(default_store="main")
|
||||
self.assertEqual(manager.default_store_id, "main")
|
||||
self.assertEqual(manager.stores, {})
|
||||
|
||||
def test_register_store(self):
|
||||
manager = TripletManager()
|
||||
store = manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
self.assertIsInstance(store, TripletStore)
|
||||
self.assertEqual(store.store_id, "main")
|
||||
self.assertEqual(store.store_type, "blazegraph")
|
||||
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||
def test_triplet_store_init(self, mock_blazegraph_store):
|
||||
store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999")
|
||||
self.assertEqual(store.backend_type, "blazegraph")
|
||||
self.assertEqual(store.endpoint, "http://localhost:9999")
|
||||
self.assertIn("main", manager.stores)
|
||||
mock_blazegraph_store.assert_called_once()
|
||||
|
||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
||||
def test_add_triplet(self, mock_get_store_backend):
|
||||
manager = TripletManager()
|
||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
|
||||
mock_store = MagicMock()
|
||||
mock_get_store_backend.return_value = mock_store
|
||||
mock_store.add_triplet.return_value = {"status": "success"}
|
||||
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||
def test_add_triplet(self, mock_blazegraph_store):
|
||||
# Setup mock backend
|
||||
mock_backend_instance = MagicMock()
|
||||
mock_blazegraph_store.return_value = mock_backend_instance
|
||||
mock_backend_instance.add_triplet.return_value = {"status": "success"}
|
||||
|
||||
store = TripletStore(backend="blazegraph")
|
||||
triplet = Triplet(subject="s", predicate="p", object="o")
|
||||
result = manager.add_triplet(triplet, store_id="main")
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(result["store_id"], "main")
|
||||
mock_store.add_triplet.assert_called_once_with(triplet)
|
||||
result = store.add_triplet(triplet)
|
||||
|
||||
self.assertEqual(result, {"status": "success"})
|
||||
mock_backend_instance.add_triplet.assert_called_once_with(triplet)
|
||||
|
||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
||||
def test_add_triplets(self, mock_get_store_backend):
|
||||
manager = TripletManager()
|
||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||
def test_add_triplets(self, mock_blazegraph_store):
|
||||
# Setup mock backend and bulk loader
|
||||
mock_backend_instance = MagicMock()
|
||||
mock_blazegraph_store.return_value = mock_backend_instance
|
||||
|
||||
mock_store = MagicMock()
|
||||
mock_get_store_backend.return_value = mock_store
|
||||
mock_store.add_triplets.return_value = {"status": "success"}
|
||||
store = TripletStore(backend="blazegraph")
|
||||
|
||||
# Mock bulk loader
|
||||
mock_loader = MagicMock()
|
||||
store.bulk_loader = mock_loader
|
||||
mock_progress = MagicMock()
|
||||
mock_progress.metadata = {"success": True}
|
||||
mock_progress.total_triplets = 2
|
||||
mock_progress.loaded_triplets = 2
|
||||
mock_progress.failed_triplets = 0
|
||||
mock_progress.total_batches = 1
|
||||
mock_loader.load_triplets.return_value = mock_progress
|
||||
|
||||
triplets = [
|
||||
Triplet(subject="s1", predicate="p1", object="o1"),
|
||||
Triplet(subject="s2", predicate="p2", object="o2")
|
||||
]
|
||||
|
||||
result = manager.add_triplets(triplets, store_id="main", batch_size=2)
|
||||
result = store.add_triplets(triplets, batch_size=2)
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(result["store_id"], "main")
|
||||
mock_store.add_triplets.assert_called()
|
||||
self.assertEqual(result["total"], 2)
|
||||
mock_loader.load_triplets.assert_called_once()
|
||||
|
||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
||||
def test_get_triplets(self, mock_get_store_backend):
|
||||
manager = TripletManager()
|
||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
|
||||
mock_store = MagicMock()
|
||||
mock_get_store_backend.return_value = mock_store
|
||||
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||
def test_get_triplets(self, mock_blazegraph_store):
|
||||
mock_backend_instance = MagicMock()
|
||||
mock_blazegraph_store.return_value = mock_backend_instance
|
||||
expected_triplets = [Triplet(subject="s", predicate="p", object="o")]
|
||||
mock_store.get_triplets.return_value = expected_triplets
|
||||
mock_backend_instance.get_triplets.return_value = expected_triplets
|
||||
|
||||
result = manager.get_triplets(subject="s", store_id="main")
|
||||
store = TripletStore(backend="blazegraph")
|
||||
result = store.get_triplets(subject="s")
|
||||
|
||||
self.assertEqual(result, expected_triplets)
|
||||
mock_store.get_triplets.assert_called_once_with("s", None, None)
|
||||
mock_backend_instance.get_triplets.assert_called_once_with(subject="s", predicate=None, object=None)
|
||||
|
||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
||||
def test_delete_triplet(self, mock_get_store_backend):
|
||||
manager = TripletManager()
|
||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
|
||||
mock_store = MagicMock()
|
||||
mock_get_store_backend.return_value = mock_store
|
||||
mock_store.delete_triplet.return_value = {"status": "deleted"}
|
||||
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||
def test_delete_triplet(self, mock_blazegraph_store):
|
||||
mock_backend_instance = MagicMock()
|
||||
mock_blazegraph_store.return_value = mock_backend_instance
|
||||
mock_backend_instance.delete_triplet.return_value = {"success": True}
|
||||
|
||||
store = TripletStore(backend="blazegraph")
|
||||
triplet = Triplet(subject="s", predicate="p", object="o")
|
||||
result = manager.delete_triplet(triplet, store_id="main")
|
||||
|
||||
result = store.delete_triplet(triplet)
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
mock_store.delete_triplet.assert_called_once_with(triplet)
|
||||
|
||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
||||
def test_update_triplet(self, mock_get_store_backend):
|
||||
manager = TripletManager()
|
||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
||||
|
||||
mock_store = MagicMock()
|
||||
mock_get_store_backend.return_value = mock_store
|
||||
mock_store.delete_triplet.return_value = {"status": "deleted"}
|
||||
mock_store.add_triplet.return_value = {"status": "added"}
|
||||
|
||||
old_triplet = Triplet(subject="s", predicate="p", object="o_old")
|
||||
new_triplet = Triplet(subject="s", predicate="p", object="o_new")
|
||||
|
||||
result = manager.update_triplet(old_triplet, new_triplet, store_id="main")
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
mock_store.delete_triplet.assert_called_once_with(old_triplet)
|
||||
mock_store.add_triplet.assert_called_once_with(new_triplet)
|
||||
|
||||
def test_query_engine_init(self):
|
||||
engine = QueryEngine(enable_caching=True)
|
||||
self.assertTrue(engine.enable_caching)
|
||||
self.assertEqual(engine.query_cache, {})
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
mock_backend_instance.delete_triplet.assert_called_once_with(triplet)
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
|
||||
import unittest
|
||||
import numpy as np
|
||||
import time
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
from semantica.vector_store import VectorStore
|
||||
from semantica.utils.exceptions import ProcessingError
|
||||
|
||||
|
||||
class TestVectorStoreParallel(unittest.TestCase):
|
||||
def setUp(self):
|
||||
logging.getLogger("vector_store").setLevel(logging.ERROR)
|
||||
|
||||
self.dimension = 4
|
||||
self.store = VectorStore(
|
||||
backend="inmemory",
|
||||
dimension=self.dimension,
|
||||
)
|
||||
|
||||
self.store.embedder = MagicMock()
|
||||
|
||||
def test_embed_batch_success(self):
|
||||
texts = ["a", "b", "c"]
|
||||
expected_embeddings = [
|
||||
np.array([0.1] * 4, dtype=np.float32),
|
||||
np.array([0.2] * 4, dtype=np.float32),
|
||||
np.array([0.3] * 4, dtype=np.float32),
|
||||
]
|
||||
|
||||
self.store.embedder.generate_embeddings.return_value = expected_embeddings
|
||||
|
||||
results = self.store.embed_batch(texts)
|
||||
|
||||
self.assertEqual(len(results), 3)
|
||||
self.assertTrue(np.allclose(results[0], expected_embeddings[0]))
|
||||
self.store.embedder.generate_embeddings.assert_called_once_with(texts)
|
||||
|
||||
def test_embed_batch_fallback(self):
|
||||
texts = ["a", "b"]
|
||||
|
||||
self.store.embedder.generate_embeddings.side_effect = Exception("Model error")
|
||||
|
||||
results = self.store.embed_batch(texts)
|
||||
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertEqual(results[0].shape, (self.dimension,))
|
||||
self.assertTrue(isinstance(results[0], np.ndarray))
|
||||
|
||||
def test_add_documents_empty(self):
|
||||
ids = self.store.add_documents([])
|
||||
self.assertEqual(ids, [])
|
||||
|
||||
def test_add_documents_metadata_mismatch(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.add_documents(["doc1"], metadata=[{}, {}])
|
||||
|
||||
def test_add_documents_parallel_success(self):
|
||||
num_docs = 10
|
||||
documents = [f"doc_{i}" for i in range(num_docs)]
|
||||
metadata = [{"id": i} for i in range(num_docs)]
|
||||
|
||||
def mock_embed_batch(texts):
|
||||
return [np.full(self.dimension, float(i)) for i, _ in enumerate(texts)]
|
||||
|
||||
with patch.object(self.store, "embed_batch", side_effect=mock_embed_batch):
|
||||
ids = self.store.add_documents(
|
||||
documents,
|
||||
metadata,
|
||||
batch_size=2,
|
||||
parallel=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(ids), num_docs)
|
||||
self.assertEqual(len(self.store.vectors), num_docs)
|
||||
|
||||
for i, vec_id in enumerate(ids):
|
||||
stored_meta = self.store.get_metadata(vec_id)
|
||||
self.assertEqual(stored_meta["id"], i)
|
||||
|
||||
def test_add_documents_sequential_success(self):
|
||||
num_docs = 5
|
||||
documents = [f"doc_{i}" for i in range(num_docs)]
|
||||
|
||||
with patch.object(self.store, "embed_batch") as mock_batch:
|
||||
mock_batch.return_value = [np.zeros(self.dimension) for _ in range(num_docs)]
|
||||
|
||||
ids = self.store.add_documents(documents, parallel=False)
|
||||
|
||||
self.assertEqual(len(ids), num_docs)
|
||||
self.assertEqual(mock_batch.call_count, 1)
|
||||
|
||||
def test_add_documents_error_propagation(self):
|
||||
documents = ["doc1", "doc2"]
|
||||
|
||||
with patch.object(self.store, "embed_batch", side_effect=ValueError("Embedding Error")):
|
||||
with self.assertRaises(Exception):
|
||||
self.store.add_documents(documents, parallel=True)
|
||||
|
||||
def test_performance_simulation(self):
|
||||
num_batches = 4
|
||||
batch_delay = 0.1
|
||||
batch_size = 1
|
||||
documents = [f"doc_{i}" for i in range(num_batches)]
|
||||
|
||||
def slow_embed(texts):
|
||||
time.sleep(batch_delay)
|
||||
return [np.zeros(self.dimension) for _ in texts]
|
||||
|
||||
with patch.object(self.store, "embed_batch", side_effect=slow_embed):
|
||||
start_seq = time.time()
|
||||
self.store.add_documents(documents, batch_size=batch_size, parallel=False)
|
||||
dur_seq = time.time() - start_seq
|
||||
|
||||
self.store.vectors = {}
|
||||
|
||||
start_par = time.time()
|
||||
self.store.add_documents(documents, batch_size=batch_size, parallel=True)
|
||||
dur_par = time.time() - start_par
|
||||
|
||||
print(f"\nPerformance Test:")
|
||||
print(f"Sequential Duration: {dur_seq:.4f}s")
|
||||
print(f"Parallel Duration: {dur_par:.4f}s")
|
||||
print(f"Speedup: {dur_seq / dur_par:.2f}x")
|
||||
|
||||
self.assertLess(dur_par, dur_seq * 0.7)
|
||||
|
||||
def test_add_documents_batch_size_edge_cases(self):
|
||||
documents = ["a", "b", "c"]
|
||||
|
||||
with patch.object(self.store, "embed_batch") as mock_batch:
|
||||
mock_batch.side_effect = lambda texts: [np.zeros(4) for _ in texts]
|
||||
|
||||
self.store.add_documents(documents, batch_size=100)
|
||||
self.assertEqual(mock_batch.call_count, 1)
|
||||
|
||||
mock_batch.reset_mock()
|
||||
|
||||
self.store.add_documents(documents, batch_size=1)
|
||||
self.assertEqual(mock_batch.call_count, 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user