Skip to content

Commit 1b5ebc2

Browse files
tekrajchhetritekrajchhetri
authored
Merge pull request #62 from sensein/improvement
StructSense Pipeline Upgrade: Specialized Tools, Robust Chunking, and BioPortal Integration
2 parents 255d29e + 55f5610 commit 1b5ebc2

516 files changed

Lines changed: 6734665 additions & 34872 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.env_example‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
OLLAMA_API_ENDPOINT=http://localhost:11434
2+
OLLAMA_MODEL=nomic-embed-text
3+
EXTERNAL_PDF_EXTRACTION_SERVICE=false
4+
GROBID_SERVER_URL_OR_EXTERNAL_SERVICE=http://localhost:8070
5+
BIOPORTAL_API_KEY=YOUR_API_KEY
6+
MAX_CONCEPT_MAPPING_RESULTS=2
7+
CREWAI_TRACING_ENABLED=false
8+
CREWAI_DISABLE_TELEMETRY=true
9+
CREWAI_DISABLE_TRACING=true
10+
CREWAI_TELEMETRY=false
11+
OTEL_SDK_DISABLED=true
12+
ENABLE_HUMAN_FEEDBACK=false
13+
ENABLE_CREW_MEMORY=false

‎.github/workflows/test.yaml‎

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ on:
55
branches:
66
- main
77
pull_request:
8+
workflow_dispatch:
89

910
jobs:
1011
unit:
@@ -27,13 +28,71 @@ jobs:
2728
python -m pip install poetry==2.3.2
2829
- name: Install dependencies with Poetry
2930
run: |
30-
poetry lock --no-update
31+
poetry lock
3132
poetry install --with dev
3233
shell: bash
3334
- name: Run unit tests
3435
id: run-tests
3536
run: >
3637
poetry run pytest \
38+
-m "not requires_openrouter" \
39+
--junitxml=pytest.xml \
40+
--cov-report=term-missing:skip-covered \
41+
--cov-report=xml:coverage.xml \
42+
--cov=src src/tests \
43+
--log-level=DEBUG \
44+
--verbose
45+
shell: bash
46+
- name: Upload coverage to Codecov
47+
uses: codecov/codecov-action@v4
48+
with:
49+
token: ${{ secrets.CODECOV_TOKEN }}
50+
51+
check-secrets:
52+
runs-on: ubuntu-latest
53+
outputs:
54+
has-openrouter-key: ${{ steps.check.outputs.has-key }}
55+
steps:
56+
- id: check
57+
env:
58+
KEY: ${{ secrets.OPENROUTER_API_KEY }}
59+
run: |
60+
if [ -n "$KEY" ]; then
61+
echo "has-key=true" >> $GITHUB_OUTPUT
62+
echo "OPENROUTER_API_KEY is set — integration tests will run."
63+
else
64+
echo "has-key=false" >> $GITHUB_OUTPUT
65+
echo "OPENROUTER_API_KEY is not set — skipping integration tests."
66+
fi
67+
shell: bash
68+
69+
integration:
70+
needs: check-secrets
71+
if: needs.check-secrets.outputs.has-openrouter-key == 'true'
72+
runs-on: ubuntu-latest
73+
steps:
74+
- uses: actions/checkout@v4
75+
with:
76+
fetch-depth: 1
77+
- name: Set up Python 3.11
78+
uses: actions/setup-python@v5
79+
with:
80+
python-version: '3.11'
81+
- name: Install poetry
82+
run: |
83+
python -m pip install --upgrade pip
84+
python -m pip install poetry==2.3.2
85+
- name: Install dependencies with Poetry
86+
run: |
87+
poetry lock
88+
poetry install --with dev
89+
shell: bash
90+
- name: Run OpenRouter integration tests
91+
env:
92+
OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
93+
run: >
94+
poetry run pytest \
95+
-m "requires_openrouter" \
3796
--junitxml=pytest.xml \
3897
--cov-report=term-missing:skip-covered \
3998
--cov-report=xml:coverage.xml \
@@ -66,7 +125,7 @@ jobs:
66125
- name: Install dependencies with Poetry
67126
run: |
68127
poetry env use ${{ matrix.python-version }}
69-
poetry lock --no-update
128+
poetry lock
70129
poetry install --with dev
71130
poetry env info
72131
shell: bash

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ example/.structsense_env
1717
.chainlit/
1818
# Distribution / packaging
1919
.Python
20+
internal_docs/
2021
build/
2122
develop-eggs/
2223
dist/

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ repos:
5151
hooks:
5252
- id: codespell
5353
args:
54-
- --skip=poetry.lock,docs_style/pdoc-theme/syntax-highlighting.css,*.cha,*.ipynb,example/sample_metadata.csv
54+
- --skip=poetry.lock,docs_style/pdoc-theme/syntax-highlighting.css,*.cha,*.ipynb,example/sample_metadata.csv,evaluation/benchmark,evaluation/pdf2reproschema/old,evaluation/ner/old,evaluation/pdf2reproschema/old/evaluation
5555
- --ignore-words-list=SIE,sie
5656

5757
- repo: https://github.com/hija/clean-dotenv

‎CHANGE_LOG.md‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# What's Changed?
2+
This document highlights the major changes made to the `StructSense`.
3+
## New Changes
4+
5+
---
6+
### 1. Implemented a new specialized `NER tool` for Named Entity Recognition to improve accuracy.
7+
The current version implements the following specialized tool (see figure) for NER task using both LLM based approach and domain specific models.
8+
![](docs/design_docs/ner_tool.png)
9+
### 2. Implements the specialized concept mapping tool, `Concept Mapping tool`, avoiding the need for manually maintaining the ontology database.
10+
- The concept mapping tool implements features such as throttling (see default values below), i.e., minimum seconds between BioPortal requests (throttle; lower = faster, higher 429 risk). This done to take into account the rate limit of the bioportal.
11+
12+
| Variable | Where used | Default | Description |
13+
|----------|------------|---------|-------------|
14+
| `BIOPORTAL_API_KEY` | `conceptmappingtool.py` | — | **Required** for ConceptMappingTool. Used in Authorization header. |
15+
| `BIOPORTAL_REQUEST_INTERVAL` | `conceptmappingtool.py` | `0.7` | Min seconds between BioPortal requests (throttle; lower = faster, higher 429 risk). |
16+
| `BIOPORTAL_BACKOFF_AFTER_429` | `conceptmappingtool.py` | `2.0` | Backoff seconds after 429. |
17+
| `MAX_CONCEPT_MAPPING_RESULTS` | `conceptmappingtool.py` | `1` | Top-N results per concept (1–50). |
18+
| `CONCEPT_MAPPING_CACHE_SIZE` | `conceptmappingtool.py` | `2000` | Max in-memory cache entries for concept mapping. |
19+
| `CONCEPT_MAPPING_MAX_TERMS` | `postprocessing.py` | — | Optional cap on how many unique terms are mapped in batch (rest get null). |
20+
21+
- To avoid the continuous API call, it implements the inmemory caching (to be replaced with new caching mechanism), i.e., `_CONCEPT_MAPPING_CACHE` (dict). The cache size is default `2000` set with `CONCEPT_MAPPING_CACHE_SIZE`. The FIFO-style cache eviction is applied, i.e., when the cache is full, one oldest entry is removed before inserting a new one. The lookup is done using `_map_single_concept()` before making an API call.
22+
### 3. Implements automatic task detection
23+
This version of **`StructSense`** introduces automatic task detection to support accurate task-to-tool mapping, enabling agents to invoke the most relevant tools for a given task.
24+
25+
The implementation combines an **LLM-based classifier** with a **heuristic fallback (e.g., see table below)** that relies on keyword matching from the task description, ensuring robustness when the LLM signal is weak or unavailable.
26+
| Condition (keywords in description) | Returned task type |
27+
|-------------------------------------|--------------------|
28+
| `"ner"` or `"named entity"` or `"entity"` | `ner` |
29+
| `"resource"` and (`"extract"` or `"dataset"` or `"tool"` or `"model"`) | `resource` |
30+
| `"structured extraction"` or `"structured_extraction"` | `structured_extraction` |
31+
| Otherwise | `extraction` |
32+
33+
Task detection is executed **once per stage** at the point where tools and post-processing are selected (e.g., `information_extraction_task`).
34+
In the full pipeline, the **first stage’s task type** determines the post-processing key and the default result shape. While each stage may theoretically define its own task type, the initial stage serves as the canonical reference for downstream processing.
35+
36+
### 4. Fixed chunking behavior
37+
The current implementation resolves previous issues with **chunking**. In earlier versions, chunking did not function correctly, which caused extraction and downstream processing to behave unexpectedly when handling chunked inputs.
38+
39+
### 5. Task-dependent post-processing
40+
41+
This version introduces task-specific post-processing logic, such as merging results across multiple chunks and globalizing chunk-level outputs into a unified result.
42+
These enhancements, e.g., `ner_post_process` `merge_ner_results` for NER are implemented in `postprocessing.py`.
43+
44+
### 6. Agent context manager
45+
This version of **StructSense** also introduces a **context manager** to support tasks such as **token compression**, enabling more controlled and efficient handling of contextual information during processing.
46+
47+
- **File:** `src/utils/context_window_manager.py`
48+
Handles token budgeting and compression for downstream agents.
49+
- **Key Components:**
50+
- `ContextWindowManager`
51+
- `prepare_for_downstream_agent`
52+
- `estimate_tokens
53+
54+
- **File:** `src/utils/downstream_agent_helper.py`
55+
Prepares inputs for alignment, judge, and human feedback agents, and merges results produced across multiple chunks.
56+
57+
- **Key Components:**
58+
- `prepare_alignment_agent_input`
59+
- `prepare_judge_agent_input`
60+
- `prepare_humanfeedback_agent_input`
61+
- `merge_structured_chunk_results`
62+
63+
- **File:** `src/utils/agent_context.py`
64+
Maintains per-agent results, state, and confidence information in a thread-safe manner.
65+
- **Key Components:**
66+
- `AgentContext`
67+
- `ThreadSafeMemory`
68+
69+
### Partial Runs
70+
71+
This version of **StructSense** supports **partial execution**, allowing you to run individual agents independently instead of executing the full pipeline.
72+
73+
For detailed examples and usage instructions, see the **[tutorial](tutorial/)** directory.
74+
75+
76+
---
77+
## Legacy
78+
### . No longer using Weaviate vector database
79+
We are no longer using the Weaviate vector database; therefore, the following keys are not required at this time. The utility code and corresponding environment variables remain in the codebase due to planned future use.
80+
#### [Weaviate](https://weaviate.io/) Configuration
81+
This configuration is optional and only necessary if you plan to integrate a knowledge source (e.g., a vector store) into the pipeline.
82+
83+
| Variable | Description | Default |
84+
|---------------------------|----------------------------------------------|-----------|
85+
| `WEAVIATE_HTTP_HOST` | HTTP host for Weaviate | `localhost` |
86+
| `WEAVIATE_HTTP_PORT` | HTTP port for Weaviate | `8080` |
87+
| `WEAVIATE_HTTP_SECURE` | Use HTTPS for HTTP connection (`true/false`) | `false` |
88+
| `WEAVIATE_GRPC_HOST` | gRPC host for Weaviate | `localhost` |
89+
| `WEAVIATE_GRPC_PORT` | gRPC port for Weaviate | `50051` |
90+
| `WEAVIATE_GRPC_SECURE` | Use secure gRPC (`true/false`) | `false` |
91+
92+
#### 🧪 Weaviate Timeouts
93+
94+
| Variable | Description | Default |
95+
|---------------------------|----------------------------------------------|-----------|
96+
| `WEAVIATE_TIMEOUT_INIT` | Timeout for initialization (in seconds) | `30` |
97+
| `WEAVIATE_TIMEOUT_QUERY` | Timeout for query operations (in seconds) | `60` |
98+
| `WEAVIATE_TIMEOUT_INSERT` | Timeout for data insertions (in seconds) | `120` |
99+
100+
#### 🤖 Ollama Configuration for WEAVIATE
101+
102+
| Variable | Description | Default |
103+
|-----------------------|-----------------------------------------------|-----------------------------------------|
104+
| `OLLAMA_API_ENDPOINT` | API endpoint for Ollama model | `http://host.docker.internal:11434` |
105+
| `OLLAMA_MODEL` | Name of the Ollama embedding model | `nomic-embed-text` |
106+
107+
> **Note**: If ollama is running in host machine and vector database, i.e., WEAVIATE, in docker, then we use `http://host.docker.internal:11434`, which is also the default value. However, if both are running in docker in the same host, use `http://localhost:11434 `.
108+
#### 🧵 Optional: Experiment Tracking
109+
110+
| Variable | Description | Default |
111+
|------------------------|----------------------------------------------------------------------------|-------------------|
112+
| `ENABLE_WEIGHTSANDBIAS` | Enable [Weights & Biases](https://wandb.ai/site) monitoring (`true/false`) | `false` |
113+
| `ENABLE_MLFLOW` | Enable [MLflow](https://mlflow.org/) logging (`true/false`) | `false` |
114+
| `MLFLOW_TRACKING_URL` | MLflow tracking server URL | `http://localhost:5000` |
115+
> **Note**: `WEAVIATE_API_KEY` is only required when `ENABLE_KG_SOURCE=true`. For Weights & Biases you need to create a project and provide its key.
116+
---

0 commit comments

Comments
 (0)