omop-llm 0.1.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,43 @@
1
+ name: docs
2
+ on:
3
+ workflow_run:
4
+ workflows: ["Release"]
5
+ types:
6
+ - completed
7
+
8
+ permissions:
9
+ contents: write
10
+
11
+ jobs:
12
+ deploy:
13
+ if: ${{ github.event.workflow_run.conclusion == 'success' }}
14
+ runs-on: ubuntu-latest
15
+ steps:
16
+ - uses: actions/checkout@v4
17
+
18
+ - name: Configure Git Credentials
19
+ run: |
20
+ git config user.name github-actions[bot]
21
+ git config user.email 41898282+github-actions[bot]@users.noreply.github.com
22
+
23
+ - uses: actions/setup-python@v5
24
+ with:
25
+ python-version: 3.x
26
+
27
+ - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
28
+
29
+ - uses: actions/cache@v4
30
+ with:
31
+ key: mkdocs-material-${{ env.cache_id }}
32
+ path: ~/.cache
33
+ restore-keys: |
34
+ mkdocs-material-
35
+
36
+ - name: Install MkDocs and Plugins
37
+ run: pip install mkdocs mkdocs-material mkdocstrings-python mkdocs-mermaid2-plugin mkdocs-table-reader-plugin
38
+
39
+ - name: Install omop-llm
40
+ run: pip install .
41
+
42
+ - name: Deploy Docs
43
+ run: mkdocs gh-deploy --force
@@ -0,0 +1,17 @@
1
+ # This workflow runs a linter on pull requests to ensure the right prefix is used for semantic-versioning
2
+ name: "Lint PR"
3
+
4
+ on:
5
+ pull_request:
6
+ types:
7
+ - opened
8
+ - edited
9
+ - synchronize
10
+
11
+ jobs:
12
+ main:
13
+ runs-on: ubuntu-latest
14
+ steps:
15
+ - uses: amannn/action-semantic-pull-request@v5
16
+ env:
17
+ GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -0,0 +1,67 @@
1
+ name: Release
2
+
3
+ on:
4
+ push:
5
+ branches:
6
+ - main
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ubuntu-latest
11
+
12
+ steps:
13
+ - uses: actions/checkout@v5
14
+
15
+ - name: Install uv
16
+ uses: astral-sh/setup-uv@v7
17
+ with:
18
+ enable-cache: true
19
+ python-version: "3.12"
20
+
21
+ - name: Install dependencies
22
+ run: uv sync --all-extras --dev
23
+
24
+ - name: Run tests
25
+ run: uv run pytest
26
+
27
+ release:
28
+ needs: test
29
+ runs-on: ubuntu-latest
30
+ concurrency: release
31
+ permissions:
32
+ id-token: write # Required for Trusted Publishing to PyPI
33
+ contents: write # Required for Semantic Release to push tags/labels
34
+
35
+ environment:
36
+ name: pypi
37
+ url: https://pypi.org/p/omop-llm
38
+
39
+ steps:
40
+ - name: Checkout
41
+ uses: actions/checkout@v5
42
+ with:
43
+ fetch-depth: 0
44
+ token: ${{ secrets.GITHUB_TOKEN }}
45
+
46
+ - name: Semantic Release
47
+ id: semantic
48
+ uses: cycjimmy/semantic-release-action@v6
49
+ with:
50
+ extra_plugins: |
51
+ @semantic-release/changelog
52
+ @semantic-release/git
53
+ @semantic-release/exec
54
+ env:
55
+ GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
56
+
57
+ - name: Install uv
58
+ if: steps.semantic.outputs.new_release_published == 'true'
59
+ uses: astral-sh/setup-uv@v7
60
+ with:
61
+ python-version: "3.12"
62
+
63
+ - name: Build and publish to PyPI
64
+ if: steps.semantic.outputs.new_release_published == 'true'
65
+ run: |
66
+ uv build
67
+ uv publish
@@ -0,0 +1,207 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ #Pipfile.lock
96
+
97
+ # UV
98
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ #uv.lock
102
+
103
+ # poetry
104
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
105
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
106
+ # commonly ignored for libraries.
107
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
108
+ #poetry.lock
109
+ #poetry.toml
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
114
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
115
+ #pdm.lock
116
+ #pdm.toml
117
+ .pdm-python
118
+ .pdm-build/
119
+
120
+ # pixi
121
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
122
+ #pixi.lock
123
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
124
+ # in the .venv directory. It is recommended not to include this directory in version control.
125
+ .pixi
126
+
127
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
128
+ __pypackages__/
129
+
130
+ # Celery stuff
131
+ celerybeat-schedule
132
+ celerybeat.pid
133
+
134
+ # SageMath parsed files
135
+ *.sage.py
136
+
137
+ # Environments
138
+ .env
139
+ .envrc
140
+ .venv
141
+ env/
142
+ venv/
143
+ ENV/
144
+ env.bak/
145
+ venv.bak/
146
+
147
+ # Spyder project settings
148
+ .spyderproject
149
+ .spyproject
150
+
151
+ # Rope project settings
152
+ .ropeproject
153
+
154
+ # mkdocs documentation
155
+ /site
156
+
157
+ # mypy
158
+ .mypy_cache/
159
+ .dmypy.json
160
+ dmypy.json
161
+
162
+ # Pyre type checker
163
+ .pyre/
164
+
165
+ # pytype static type analyzer
166
+ .pytype/
167
+
168
+ # Cython debug symbols
169
+ cython_debug/
170
+
171
+ # PyCharm
172
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
173
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
174
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
175
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
176
+ #.idea/
177
+
178
+ # Abstra
179
+ # Abstra is an AI-powered process automation framework.
180
+ # Ignore directories containing user credentials, local state, and settings.
181
+ # Learn more at https://abstra.io/docs
182
+ .abstra/
183
+
184
+ # Visual Studio Code
185
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
186
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
187
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
188
+ # you could uncomment the following to ignore the entire vscode folder
189
+ # .vscode/
190
+
191
+ # Ruff stuff:
192
+ .ruff_cache/
193
+
194
+ # PyPI configuration file
195
+ .pypirc
196
+
197
+ # Cursor
198
+ # Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
199
+ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
200
+ # refer to https://docs.cursor.com/context/ignore-files
201
+ .cursorignore
202
+ .cursorindexingignore
203
+
204
+ # Marimo
205
+ marimo/_static/
206
+ marimo/_lsp/
207
+ __marimo__/
@@ -0,0 +1,22 @@
1
+ {
2
+ "branches": ["main"],
3
+ "plugins": [
4
+ "@semantic-release/commit-analyzer",
5
+ "@semantic-release/release-notes-generator",
6
+ "@semantic-release/changelog",
7
+ [
8
+ "@semantic-release/exec",
9
+ {
10
+ "prepareCmd": "sed -i 's/version = \".*\"/version = \"${nextRelease.version}\"/' pyproject.toml"
11
+ }
12
+ ],
13
+ [
14
+ "@semantic-release/git",
15
+ {
16
+ "assets": ["CHANGELOG.md", "pyproject.toml"],
17
+ "message": "chore(release): ${nextRelease.version} [skip ci]\n\n${nextRelease.notes}"
18
+ }
19
+ ],
20
+ "@semantic-release/github"
21
+ ]
22
+ }
@@ -0,0 +1,15 @@
1
+ ## [0.1.1](https://github.com/AustralianCancerDataNetwork/omop-llm/compare/v0.1.0...v0.1.1) (2026-04-01)
2
+
3
+
4
+ ### Bug Fixes
5
+
6
+ * trigger initial 0.x release ([63ce049](https://github.com/AustralianCancerDataNetwork/omop-llm/commit/63ce04934bdb021b903b87dfd452f2e028324bd0))
7
+
8
+ # 1.0.0 (2026-04-01)
9
+
10
+
11
+ ### Bug Fixes
12
+
13
+ * Accurate dependencies for deployment ([5ebdfb2](https://github.com/AustralianCancerDataNetwork/omop-llm/commit/5ebdfb2ba088f2dda9cbd0c554c6f4a52fc119e5))
14
+ * direct deployment to pypi ([33bef1a](https://github.com/AustralianCancerDataNetwork/omop-llm/commit/33bef1a2f63493c0d0e3dd71a40d3f214302150f))
15
+ * include releaserc for semantic versioning ([fe62f2f](https://github.com/AustralianCancerDataNetwork/omop-llm/commit/fe62f2f421e9fda61d13bd7be39b14df9b8b816b))
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.4
2
+ Name: omop-llm
3
+ Version: 0.1.1
4
+ Summary: LLM interfaces for OMOP
5
+ Author-email: Nico Loesch <n.loesch@unsw.edu.au>
6
+ License-Expression: Apache-2.0
7
+ Requires-Python: >=3.12
8
+ Requires-Dist: instructor>=1.13.0
9
+ Requires-Dist: numpy
10
+ Requires-Dist: openai
11
+ Requires-Dist: prompt-spec>=0.1.4
12
+ Requires-Dist: pydantic
13
+ Description-Content-Type: text/markdown
14
+
15
+ # omop-llm
16
+ LLM interface for embedding and instructor tasks. Extended documentation can be found [here](https://AustralianCancerDataNetwork.github.io/omop-llm).
@@ -0,0 +1,2 @@
1
+ # omop-llm
2
+ LLM interface for embedding and instructor tasks. Extended documentation can be found [here](https://AustralianCancerDataNetwork.github.io/omop-llm).
@@ -0,0 +1,19 @@
1
+ # OMOP LLM Interface
2
+
3
+ `omop-llm` is designed as a simple LLM interface in the context of the OMOP CDM. Particularly, the wrapper currently exposes two interfaces:
4
+
5
+ !!! warning
6
+
7
+ The backend is realised as an OpenAI client that would support a wide variety of models but we currently only support Ollama. Extension for this is planned in future releases.
8
+
9
+ - **`LLMClient`**: Base client with the capacity to:
10
+ - obtain metadata information
11
+ - calculate embeddings on demand
12
+ - calculate semantic similarity
13
+ - **`InstructorClient`**: Child client of `LLMClient` to:
14
+ - provide an interface for the [`instructor`](https://python.useinstructor.com/) library with easy instantiation.
15
+ - chat completions using chat messages.
16
+
17
+
18
+ ## Documentation overview
19
+ - [Installation](usage/installation.md)
@@ -0,0 +1,10 @@
1
+ # API Reference
2
+
3
+ This reference is automatically generated from the source code.
4
+
5
+ ::: omop_llm
6
+ options:
7
+ members:
8
+ - interface
9
+ show_submodules: true
10
+ show_root_heading: true
@@ -0,0 +1,7 @@
1
+ # Installation instructions
2
+
3
+ The package is currently only available from GitHub.
4
+
5
+ ```bash
6
+ pip install git+https://github.com/AustralianCancerDataNetwork/omop-llm.git
7
+ ```
@@ -0,0 +1,51 @@
1
+ site_name: OMOP LLM
2
+ site_description: LLM interface for the OMOP Graph project for easy LLM tooling.
3
+ site_url: https://AustralianCancerDataNetwork.github.io/omop-llm
4
+ repo_url: https://github.com/AustralianCancerDataNetwork/omop-llm
5
+ repo_name: omop-llm
6
+
7
+ theme:
8
+ name: material
9
+ features:
10
+ - navigation.sections
11
+ - navigation.expand
12
+ - navigation.path
13
+ - navigation.top
14
+ - content.code.copy
15
+ - content.code.annotate
16
+ palette:
17
+ - scheme: default
18
+ primary: indigo
19
+ accent: blue
20
+ toggle:
21
+ icon: material/brightness-7
22
+ name: Switch to dark mode
23
+ - scheme: slate
24
+ primary: indigo
25
+ accent: light-blue
26
+ toggle:
27
+ icon: material/brightness-4
28
+ name: Switch to light mode
29
+
30
+ nav:
31
+ - Home: index.md
32
+ - Installation: usage/installation.md
33
+ - "API Reference": reference.md
34
+
35
+ plugins:
36
+ - search
37
+ - mkdocstrings:
38
+ handlers:
39
+ python:
40
+ paths: [src]
41
+ options:
42
+ docstring_style: numpy
43
+ show_root_heading: true
44
+ show_source: false
45
+
46
+ markdown_extensions:
47
+ - admonition
48
+ - pymdownx.details
49
+ - pymdownx.superfences
50
+ - attr_list
51
+ - tables
@@ -0,0 +1,47 @@
1
+ [project]
2
+ name = "omop-llm"
3
+ version = "0.1.1"
4
+ description = "LLM interfaces for OMOP"
5
+ authors = [
6
+ {name = "Nico Loesch", email = "n.loesch@unsw.edu.au"}
7
+ ]
8
+ license = "Apache-2.0"
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ dependencies = [
12
+ "instructor>=1.13.0",
13
+ "prompt-spec>=0.1.4",
14
+ "openai",
15
+ "numpy",
16
+ "pydantic",
17
+ ]
18
+
19
+ [dependency-groups]
20
+ dev = [
21
+ "mypy>=1.19.1",
22
+ "pytest>=9.0.2",
23
+ "pytest-cov>=7.0.0",
24
+ "types-pyyaml>=6.0.12.20250915",
25
+ "mkdocs<2.0",
26
+ "mkdocs-material",
27
+ "mkdocstrings[python]",
28
+ "mkdocs-mermaid2-plugin"
29
+ ]
30
+
31
+ [build-system]
32
+ requires = ["hatchling"]
33
+ build-backend = "hatchling.build"
34
+
35
+ [tool.hatch.build.targets.wheel]
36
+ packages = ["src/omop_llm"]
37
+
38
+ [tool.pytest.ini_options]
39
+ testpaths = ["tests"]
40
+ python_files = ["test_*.py"]
41
+ python_classes = ["Test*"]
42
+ python_functions = ["test_*"]
43
+ addopts = "-v --tb=short"
44
+ markers = [
45
+ "unit: Unit tests that do not require external API calls",
46
+ "integration: Tests that hit live LLM endpoints (requires API keys)",
47
+ ]
@@ -0,0 +1,3 @@
1
+ from .interface import LLMClient, InstructorClient, CHAT_MESSAGE_DICT
2
+
3
+ __all__ = ["LLMClient", "InstructorClient", "CHAT_MESSAGE_DICT"]
@@ -0,0 +1,2 @@
1
+ from .client import LLMClient, CHAT_MESSAGE_DICT
2
+ from .instructor_client import InstructorClient
@@ -0,0 +1,277 @@
1
+ import logging
2
+ from dataclasses import dataclass, field
3
+ from typing import Any, Dict, List, Optional, Union, TypeAlias, Tuple
4
+
5
+ import numpy as np
6
+ import requests
7
+ from openai import OpenAI
8
+
9
+ logger = logging.getLogger(__name__)
10
+
11
+
12
+ CHAT_MESSAGE_DICT: TypeAlias = Dict[str, str]
13
+
14
+
15
+ class LLMClientError(RuntimeError):
16
+ """
17
+ Custom exception for LLM Client runtime errors.
18
+ """
19
+ pass
20
+
21
+
22
+ @dataclass
23
+ class LLMClient:
24
+ """
25
+ Base class for LLM clients.
26
+
27
+ This class replicates the API of the OntoGPT LLMClient but serves as a
28
+ base for other implementations (e.g., InstructorClient). Relies on the
29
+ OpenAI client for core functionality.
30
+
31
+ Parameters
32
+ ----------
33
+ model : str
34
+ The name of the model to use (e.g., 'gpt-4', 'llama3').
35
+ api_base : str
36
+ The base URL for the API endpoint.
37
+ api_key : str
38
+ The API key for authentication.
39
+ temperature : float, optional
40
+ The temperature parameter for generation. Default is 1.0.
41
+ system_message : str, optional
42
+ The default system message to prepend to chats. Default is "".
43
+
44
+ Attributes
45
+ ----------
46
+ _base_client : OpenAI
47
+ The initialized OpenAI client instance.
48
+ _embedding_dim : int or None
49
+ Cached embedding dimension size.
50
+ """
51
+
52
+ model: str
53
+ api_base: str
54
+ api_key: str = "ollama" # required by OpenAI client, ignored by Ollama
55
+ temperature: float = 1.0
56
+ system_message: str = ""
57
+ embedding_batch_size: int = 32
58
+ _base_client: OpenAI = field(init=False, repr=False)
59
+ _embedding_dim: Optional[int] = field(init=False, default=None)
60
+
61
+ def __post_init__(self) -> None:
62
+ logger.info(f"Initialising {self.__class__.__name__} for model={self.model}")
63
+ if self.api_key is None:
64
+ self.api_key = "ollama" # Default to "ollama" for compatibility, but can be overridden
65
+
66
+ self._base_client = OpenAI(
67
+ base_url=self.api_base,
68
+ api_key=self.api_key,
69
+ )
70
+
71
+ @property
72
+ def embedding_dim(self) -> int:
73
+ """
74
+ Retrieve the embedding dimension for the current model.
75
+
76
+ If the dimension is not cached, it attempts to fetch it from the API.
77
+ Currently supports Ollama endpoints.
78
+
79
+ Returns
80
+ -------
81
+ int
82
+ The size of the embedding vector.
83
+
84
+ Raises
85
+ ------
86
+ ValueError
87
+ If model information cannot be found in the Ollama response.
88
+ NotImplementedError
89
+ If the API base is not supported for automatic dimension retrieval.
90
+ """
91
+ if self._embedding_dim is not None:
92
+ return self._embedding_dim
93
+
94
+ if (
95
+ "ollama" in self.api_base or
96
+ (
97
+ (
98
+ "localhost" in self.api_base or
99
+ "127.0.0.1" in self.api_base
100
+ ) and self.api_key == "ollama"
101
+ )
102
+ ):
103
+ # Strip /v1 to access base Ollama API
104
+ ollama_url_without_v1 = self.api_base.replace("/v1", "")
105
+ requests_url = f"{ollama_url_without_v1}/api/show"
106
+
107
+ response = requests.post(requests_url, json={"name": self.model}).json()
108
+ model_info = response.get("model_info", {})
109
+
110
+ if model_info:
111
+ # Find keys resembling 'embedding_length'
112
+ embedding_key = [key for key in model_info.keys() if "embedding_length" in key]
113
+ if len(embedding_key) == 1:
114
+ self._embedding_dim = int(model_info[embedding_key[0]])
115
+ return self._embedding_dim
116
+
117
+ raise ValueError(f"Model information not found in Ollama response: {response}")
118
+ else:
119
+ raise NotImplementedError("Embedding dimension retrieval not implemented for this API base")
120
+
121
+ @property
122
+ def base_client(self) -> OpenAI:
123
+ return self._base_client
124
+
125
+ def embeddings(self, text: Union[str, List[str], Tuple[str, ...]], batch_size: Optional[int] = None) -> np.ndarray:
126
+ """
127
+ Retrieve embeddings for the given text.
128
+
129
+ Parameters
130
+ ----------
131
+ text : str or List[str]
132
+ The input text or list of texts to embed.
133
+ batch_size : int, optional
134
+ The number of texts to process in a single API call. Default is 32.
135
+
136
+ Returns
137
+ -------
138
+ np.ndarray
139
+ A 2D numpy array containing the embeddings.
140
+
141
+ Raises
142
+ ------
143
+ AssertionError
144
+ If the base client has not been initialized.
145
+ """
146
+ assert self.base_client is not None, "Base client should be initialized"
147
+ if batch_size is None:
148
+ batch_size = self.embedding_batch_size
149
+
150
+ if isinstance(text, str):
151
+ text = (text, )
152
+ elif isinstance(text, list):
153
+ text = tuple(text)
154
+
155
+ batch_buffer = []
156
+
157
+ for batch_chunk_idx in range(0, len(text), batch_size):
158
+ logger.debug(f"Processing batch chunk from index {batch_chunk_idx} to {batch_chunk_idx + batch_size}")
159
+ batch_chunk = text[batch_chunk_idx:batch_chunk_idx + batch_size]
160
+ response = self.base_client.embeddings.create(
161
+ model=self.model,
162
+ input=batch_chunk,
163
+ )
164
+ batch_buffer.extend([emb.embedding for emb in response.data])
165
+
166
+ return np.array(batch_buffer)
167
+
168
+ def similarity(
169
+ self,
170
+ terms: Union[str, List[str], np.ndarray],
171
+ terms_to_match: Union[str, List[str], np.ndarray],
172
+ **kwargs: Any
173
+ ) -> np.ndarray:
174
+ """
175
+ Calculate the cosine similarity between two sets of terms.
176
+
177
+ This method handles inputs as strings, lists of strings, or pre-computed
178
+ numpy arrays of embeddings.
179
+
180
+ Parameters
181
+ ----------
182
+ terms : str, List[str], or np.ndarray
183
+ The source terms or embeddings.
184
+ terms_to_match : str, List[str], or np.ndarray
185
+ The target terms or embeddings to match against.
186
+ **kwargs : Any
187
+ Additional arguments passed to the embedding function if embedding is required.
188
+
189
+ Returns
190
+ -------
191
+ np.ndarray
192
+ A similarity matrix.
193
+
194
+ Raises
195
+ ------
196
+ ValueError
197
+ If inputs are not strings, lists, or numpy arrays.
198
+ """
199
+ if isinstance(terms, str):
200
+ terms = [terms]
201
+ if isinstance(terms_to_match, str):
202
+ terms_to_match = [terms_to_match]
203
+
204
+ # Process source terms
205
+ if isinstance(terms, list):
206
+ terms_embeddings = self.embeddings(text=terms, **kwargs)
207
+ elif isinstance(terms, np.ndarray):
208
+ terms_embeddings = terms
209
+ else:
210
+ raise ValueError("terms must be either a string, list of strings, or numpy array")
211
+
212
+ # Process target terms
213
+ if isinstance(terms_to_match, list):
214
+ terms_to_match_embeddings = self.embeddings(text=terms_to_match, **kwargs)
215
+ elif isinstance(terms_to_match, np.ndarray):
216
+ terms_to_match_embeddings = terms_to_match
217
+ else:
218
+ raise ValueError("terms_to_match must be either a string, list of strings, or numpy array")
219
+
220
+ return self.cosine_similarity(terms_embeddings, terms_to_match_embeddings)
221
+
222
+ @staticmethod
223
+ def cosine_similarity(vecs_a: np.ndarray, vecs_b: np.ndarray) -> np.ndarray:
224
+ """
225
+ Compute the cosine similarity between two matrices of vectors.
226
+
227
+ Parameters
228
+ ----------
229
+ vecs_a : np.ndarray
230
+ A 2D array of vectors (Shape: M x D).
231
+ vecs_b : np.ndarray
232
+ A 2D array of vectors (Shape: N x D).
233
+
234
+ Returns
235
+ -------
236
+ np.ndarray
237
+ The dot product of the normalized vectors (Shape: M x N).
238
+
239
+ Notes
240
+ -----
241
+ A small epsilon (1e-10) is added to the norms to prevent division by zero.
242
+ """
243
+ assert vecs_a.ndim == 2 and vecs_b.ndim == 2, "Input vectors must be 2D arrays"
244
+
245
+ norm_a = np.linalg.norm(vecs_a, axis=1, keepdims=True)
246
+ norm_b = np.linalg.norm(vecs_b, axis=1, keepdims=True)
247
+
248
+ # Prevent division by zero
249
+ norm_a[norm_a == 0] = 1e-10
250
+ norm_b[norm_b == 0] = 1e-10
251
+
252
+ vecs_a_norm = vecs_a / norm_a
253
+ vecs_b_norm = vecs_b / norm_b
254
+
255
+ return np.dot(vecs_a_norm, vecs_b_norm.T)
256
+
257
+ def euclidean_distance(self, text1: str, text2: str, **kwargs: Any) -> float:
258
+ """
259
+ Calculate the Euclidean distance between embeddings of two texts.
260
+
261
+ Parameters
262
+ ----------
263
+ text1 : str
264
+ The first text string.
265
+ text2 : str
266
+ The second text string.
267
+ **kwargs : Any
268
+ Additional arguments passed to the embedding function.
269
+
270
+ Returns
271
+ -------
272
+ float
273
+ The Euclidean distance (L2 norm) between the two embedding vectors.
274
+ """
275
+ a1 = self.embeddings(text1, **kwargs)
276
+ a2 = self.embeddings(text2, **kwargs)
277
+ return float(np.linalg.norm(np.array(a1) - np.array(a2)))
@@ -0,0 +1,153 @@
1
+ import logging
2
+ from dataclasses import dataclass, field
3
+ from typing import Any, Dict, List, Optional, Union, TypeAlias
4
+
5
+ import instructor
6
+ from pydantic import BaseModel
7
+ from prompt_spec import PromptTemplate
8
+
9
+ from .client import LLMClient, LLMClientError
10
+
11
+ logger = logging.getLogger(__name__)
12
+
13
+
14
+ CHAT_MESSAGE_DICT: TypeAlias = Dict[str, str]
15
+
16
+
17
+ @dataclass
18
+ class InstructorClient(LLMClient):
19
+ """
20
+ LLMClient implementation backed by pydantic-instructor.
21
+
22
+ This client extends the base LLMClient to support structured outputs
23
+ via the `instructor` library.
24
+
25
+ Parameters
26
+ ----------
27
+ instructor_mode : instructor.Mode
28
+ The mode for the instructor client (e.g., JSON, TOOLS). Default is JSON.
29
+
30
+ Attributes
31
+ ----------
32
+ _client : Any
33
+ The initialized instructor client wrapper.
34
+ """
35
+
36
+ instructor_mode: instructor.Mode = instructor.Mode.JSON
37
+ _client: Any = field(init=False, repr=False)
38
+
39
+ def __post_init__(self) -> None:
40
+ """
41
+ Initialize the Instructor wrapper around the OpenAI client.
42
+ """
43
+ super().__post_init__()
44
+ self._client = instructor.from_openai(
45
+ self.base_client,
46
+ mode=self.instructor_mode,
47
+ )
48
+
49
+ def complete[T: BaseModel](
50
+ self,
51
+ messages: List[CHAT_MESSAGE_DICT],
52
+ response_model: Optional[type[T]] = None,
53
+ show_prompt: bool = False,
54
+ **kwargs: Any,
55
+ ) -> Union[str, T]:
56
+ """
57
+ Run a chat completion.
58
+
59
+ If `response_model` is provided, structured output is returned based
60
+ on the Pydantic model. Otherwise, plain text is returned.
61
+
62
+ Parameters
63
+ ----------
64
+ messages : list of dict
65
+ The list of chat messages (e.g., `[{'role': 'user', 'content': '...'}]`).
66
+ response_model : type[T], optional
67
+ A Pydantic model class (T) to structure the response.
68
+ Must be a subclass of BaseModel.
69
+ show_prompt : bool, optional
70
+ If True, logs the rendered prompt before sending. Default is False.
71
+ **kwargs : Any
72
+ Additional arguments passed to `chat.completions.create`.
73
+
74
+ Returns
75
+ -------
76
+ Union[str, T]
77
+ The response string (if no model provided) or an instance of T (the Pydantic model).
78
+
79
+ Raises
80
+ ------
81
+ LLMClientError
82
+ If the completion request fails.
83
+ """
84
+ if show_prompt:
85
+ # Note: render_prompt_messages only uses the 'messages' list.
86
+ rendered_text = self.render_prompt_messages(messages=messages)
87
+ logger.info(f"SENDING PROMPT:\n{rendered_text}")
88
+
89
+ try:
90
+ result = self._client.chat.completions.create(
91
+ model=self.model,
92
+ messages=messages,
93
+ temperature=self.temperature,
94
+ response_model=response_model,
95
+ **kwargs,
96
+ )
97
+ except Exception as e:
98
+ logger.error(f"Instructor completion failed: {e}")
99
+ raise LLMClientError("Instructor completion failed") from e
100
+
101
+ if response_model is not None:
102
+ # result is already an instance of T here
103
+ return result
104
+
105
+ return result.choices[0].message.content
106
+
107
+ # TODO: The next two messages should be in prompt_spec!
108
+ def messages_from_prompt_template(
109
+ self, prompt_template: Optional[PromptTemplate], text: str
110
+ ) -> List[Dict[str, str]]:
111
+ """
112
+ Generate a list of messages from a PromptTemplate.
113
+
114
+ Raises
115
+ ------
116
+ NotImplementedError
117
+ This method is currently dropped for Template in LinkML.
118
+ """
119
+ raise NotImplementedError("Method dropped for Template in LinkML")
120
+
121
+ # Unreachable code preserved for reference/future implementation
122
+ # messages = [{"role": "system", "content": prompt_template.system}] if self.system_message else []
123
+ # for example in prompt_template.examples:
124
+ # messages.append({"role": "user", "content": example.input})
125
+ # messages.append({"role": "assistant", "content": example.output.model_dump_json(indent=2)})
126
+ # messages.append({"role": "user", "content": text})
127
+ # return messages
128
+
129
+ def render_prompt_messages(self, messages: List[CHAT_MESSAGE_DICT]) -> str:
130
+ """
131
+ Render a list of chat messages into a single string for logging or display.
132
+
133
+ Parameters
134
+ ----------
135
+ messages : list of dict
136
+ The chat history.
137
+
138
+ Returns
139
+ -------
140
+ str
141
+ A formatted string representation of the conversation.
142
+ """
143
+ lines = []
144
+ for msg in messages:
145
+ role_label = "System:"
146
+ if msg['role'] == "user":
147
+ role_label = "Input:"
148
+ elif msg['role'] == "assistant":
149
+ role_label = "Output:"
150
+
151
+ lines.append(f"{role_label} {msg['content']}")
152
+
153
+ return "\n".join(lines)
@@ -0,0 +1,4 @@
1
+ def test_import():
2
+ """Dummy test to ensure the package can be imported."""
3
+ import omop_llm
4
+ assert True