proofrag-ai 0.1.0__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.
- proofrag_ai-0.1.0/LICENSE +0 -0
- proofrag_ai-0.1.0/PKG-INFO +48 -0
- proofrag_ai-0.1.0/README.md +0 -0
- proofrag_ai-0.1.0/pyproject.toml +128 -0
- proofrag_ai-0.1.0/setup.cfg +4 -0
- proofrag_ai-0.1.0/src/proofrag/__init__.py +0 -0
- proofrag_ai-0.1.0/src/proofrag/api.py +65 -0
- proofrag_ai-0.1.0/src/proofrag/cli.py +96 -0
- proofrag_ai-0.1.0/src/proofrag/config.py +108 -0
- proofrag_ai-0.1.0/src/proofrag/embeddings.py +53 -0
- proofrag_ai-0.1.0/src/proofrag/generator.py +275 -0
- proofrag_ai-0.1.0/src/proofrag/ingestion.py +379 -0
- proofrag_ai-0.1.0/src/proofrag/models.py +180 -0
- proofrag_ai-0.1.0/src/proofrag/okf.py +555 -0
- proofrag_ai-0.1.0/src/proofrag/pipeline.py +222 -0
- proofrag_ai-0.1.0/src/proofrag/retrieval.py +350 -0
- proofrag_ai-0.1.0/src/proofrag/verifier.py +272 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/PKG-INFO +48 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/SOURCES.txt +29 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/dependency_links.txt +1 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/entry_points.txt +2 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/requires.txt +28 -0
- proofrag_ai-0.1.0/src/proofrag_ai.egg-info/top_level.txt +1 -0
- proofrag_ai-0.1.0/tests/test_api.py +0 -0
- proofrag_ai-0.1.0/tests/test_generator.py +231 -0
- proofrag_ai-0.1.0/tests/test_ingestion.py +152 -0
- proofrag_ai-0.1.0/tests/test_models.py +0 -0
- proofrag_ai-0.1.0/tests/test_okf.py +307 -0
- proofrag_ai-0.1.0/tests/test_pipeline.py +344 -0
- proofrag_ai-0.1.0/tests/test_retrieval.py +162 -0
- proofrag_ai-0.1.0/tests/test_verifier.py +243 -0
|
File without changes
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: proofrag-ai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Evidence-backed RAG with verified citations and refusal on insufficient evidence.
|
|
5
|
+
Author: Mainak Sen
|
|
6
|
+
Keywords: rag,retrieval-augmented-generation,llm,citations,grounded-ai,knowledge-base,okf,chroma,langchain,generative-ai
|
|
7
|
+
Classifier: Development Status :: 3 - Alpha
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: Intended Audience :: Science/Research
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
18
|
+
Requires-Python: <3.14,>=3.10
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
License-File: LICENSE
|
|
21
|
+
Requires-Dist: pydantic<3.0,>=2.9
|
|
22
|
+
Requires-Dist: pydantic-settings<3.0,>=2.6
|
|
23
|
+
Requires-Dist: python-dotenv<2.0,>=1.0
|
|
24
|
+
Requires-Dist: PyYAML<7.0,>=6.0
|
|
25
|
+
Requires-Dist: pymupdf<2.0,>=1.24
|
|
26
|
+
Requires-Dist: langchain<2.0,>=0.3
|
|
27
|
+
Requires-Dist: langchain-community<1.0,>=0.3
|
|
28
|
+
Requires-Dist: langchain-core<2.0,>=0.3
|
|
29
|
+
Requires-Dist: langchain-chroma<1.0,>=0.2
|
|
30
|
+
Requires-Dist: langchain-huggingface<1.0,>=0.1
|
|
31
|
+
Requires-Dist: chromadb<2.0,>=0.5
|
|
32
|
+
Requires-Dist: sentence-transformers<6.0,>=3.0
|
|
33
|
+
Requires-Dist: groq<1.0,>=0.13
|
|
34
|
+
Requires-Dist: langchain-groq<1.0,>=0.2
|
|
35
|
+
Requires-Dist: fastapi<1.0,>=0.115
|
|
36
|
+
Requires-Dist: uvicorn[standard]<1.0,>=0.32
|
|
37
|
+
Requires-Dist: typer<1.0,>=0.15
|
|
38
|
+
Requires-Dist: rich<15.0,>=13.9
|
|
39
|
+
Requires-Dist: streamlit<2.0,>=1.40
|
|
40
|
+
Requires-Dist: nltk<4.0,>=3.9
|
|
41
|
+
Provides-Extra: dev
|
|
42
|
+
Requires-Dist: pytest<9.0,>=8.3; extra == "dev"
|
|
43
|
+
Requires-Dist: pytest-cov<8.0,>=6.0; extra == "dev"
|
|
44
|
+
Requires-Dist: httpx<1.0,>=0.27; extra == "dev"
|
|
45
|
+
Requires-Dist: ruff<1.0,>=0.8; extra == "dev"
|
|
46
|
+
Requires-Dist: build<2.0,>=1.2; extra == "dev"
|
|
47
|
+
Requires-Dist: twine<7.0,>=6.0; extra == "dev"
|
|
48
|
+
Dynamic: license-file
|
|
File without changes
|
|
@@ -0,0 +1,128 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=75.0", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
[project]
|
|
7
|
+
name = "proofrag-ai"
|
|
8
|
+
version = "0.1.0"
|
|
9
|
+
description = "Evidence-backed RAG with verified citations and refusal on insufficient evidence."
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10,<3.14"
|
|
12
|
+
license = { file = "LICENSE" }
|
|
13
|
+
authors = [
|
|
14
|
+
{ name = "Mainak Sen" }
|
|
15
|
+
]
|
|
16
|
+
keywords = [
|
|
17
|
+
"rag",
|
|
18
|
+
"retrieval-augmented-generation",
|
|
19
|
+
"llm",
|
|
20
|
+
"citations",
|
|
21
|
+
"grounded-ai",
|
|
22
|
+
"knowledge-base",
|
|
23
|
+
"okf",
|
|
24
|
+
"chroma",
|
|
25
|
+
"langchain",
|
|
26
|
+
"generative-ai"
|
|
27
|
+
]
|
|
28
|
+
classifiers = [
|
|
29
|
+
"Development Status :: 3 - Alpha",
|
|
30
|
+
"Intended Audience :: Developers",
|
|
31
|
+
"Intended Audience :: Science/Research",
|
|
32
|
+
"License :: OSI Approved :: MIT License",
|
|
33
|
+
"Programming Language :: Python :: 3",
|
|
34
|
+
"Programming Language :: Python :: 3.10",
|
|
35
|
+
"Programming Language :: Python :: 3.11",
|
|
36
|
+
"Programming Language :: Python :: 3.12",
|
|
37
|
+
"Programming Language :: Python :: 3.13",
|
|
38
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
39
|
+
"Topic :: Software Development :: Libraries :: Python Modules"
|
|
40
|
+
]
|
|
41
|
+
|
|
42
|
+
dependencies = [
|
|
43
|
+
"pydantic>=2.9,<3.0",
|
|
44
|
+
"pydantic-settings>=2.6,<3.0",
|
|
45
|
+
"python-dotenv>=1.0,<2.0",
|
|
46
|
+
|
|
47
|
+
"PyYAML>=6.0,<7.0",
|
|
48
|
+
|
|
49
|
+
"pymupdf>=1.24,<2.0",
|
|
50
|
+
|
|
51
|
+
"langchain>=0.3,<2.0",
|
|
52
|
+
"langchain-community>=0.3,<1.0",
|
|
53
|
+
"langchain-core>=0.3,<2.0",
|
|
54
|
+
"langchain-chroma>=0.2,<1.0",
|
|
55
|
+
"langchain-huggingface>=0.1,<1.0",
|
|
56
|
+
|
|
57
|
+
"chromadb>=0.5,<2.0",
|
|
58
|
+
"sentence-transformers>=3.0,<6.0",
|
|
59
|
+
|
|
60
|
+
"groq>=0.13,<1.0",
|
|
61
|
+
"langchain-groq>=0.2,<1.0",
|
|
62
|
+
|
|
63
|
+
"fastapi>=0.115,<1.0",
|
|
64
|
+
"uvicorn[standard]>=0.32,<1.0",
|
|
65
|
+
|
|
66
|
+
"typer>=0.15,<1.0",
|
|
67
|
+
"rich>=13.9,<15.0",
|
|
68
|
+
|
|
69
|
+
"streamlit>=1.40,<2.0",
|
|
70
|
+
"nltk>=3.9,<4.0"
|
|
71
|
+
]
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
[project.optional-dependencies]
|
|
75
|
+
dev = [
|
|
76
|
+
"pytest>=8.3,<9.0",
|
|
77
|
+
"pytest-cov>=6.0,<8.0",
|
|
78
|
+
"httpx>=0.27,<1.0",
|
|
79
|
+
"ruff>=0.8,<1.0",
|
|
80
|
+
"build>=1.2,<2.0",
|
|
81
|
+
"twine>=6.0,<7.0"
|
|
82
|
+
]
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
[project.scripts]
|
|
86
|
+
proofrag = "proofrag.cli:app"
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
[tool.setuptools]
|
|
90
|
+
package-dir = { "" = "src" }
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
[tool.setuptools.packages.find]
|
|
94
|
+
where = ["src"]
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
[tool.pytest.ini_options]
|
|
98
|
+
testpaths = ["tests"]
|
|
99
|
+
pythonpath = ["src"]
|
|
100
|
+
addopts = "-ra"
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
[tool.ruff]
|
|
104
|
+
line-length = 100
|
|
105
|
+
target-version = "py310"
|
|
106
|
+
|
|
107
|
+
[tool.ruff.lint]
|
|
108
|
+
select = [
|
|
109
|
+
"E",
|
|
110
|
+
"F",
|
|
111
|
+
"I",
|
|
112
|
+
"B",
|
|
113
|
+
"UP"
|
|
114
|
+
]
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
[tool.coverage.run]
|
|
118
|
+
source = ["proofrag"]
|
|
119
|
+
|
|
120
|
+
[tool.coverage.report]
|
|
121
|
+
exclude_lines = [
|
|
122
|
+
"pragma: no cover",
|
|
123
|
+
"if TYPE_CHECKING:"
|
|
124
|
+
]
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
[build]
|
|
128
|
+
target = "proofrag"
|
|
File without changes
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from fastapi import FastAPI, HTTPException
|
|
4
|
+
|
|
5
|
+
from .config import get_settings
|
|
6
|
+
from .models import ProofRAGRequest, ProofRAGResponse
|
|
7
|
+
from .pipeline import PipelineError, ProofRAG
|
|
8
|
+
|
|
9
|
+
settings = get_settings()
|
|
10
|
+
|
|
11
|
+
app = FastAPI(
|
|
12
|
+
title=settings.app_name,
|
|
13
|
+
version=settings.app_version,
|
|
14
|
+
description=(
|
|
15
|
+
"Evidence-backed RAG with verified citations "
|
|
16
|
+
"and safe refusal on insufficient evidence."
|
|
17
|
+
),
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
rag = ProofRAG()
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@app.get("/health")
|
|
24
|
+
def health() -> dict[str, str]:
|
|
25
|
+
return {
|
|
26
|
+
"status": "ok",
|
|
27
|
+
"service": settings.app_name,
|
|
28
|
+
"version": settings.app_version,
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@app.post("/ingest")
|
|
33
|
+
def ingest() -> dict[str, object]:
|
|
34
|
+
try:
|
|
35
|
+
documents = rag.ingest()
|
|
36
|
+
except PipelineError as exc:
|
|
37
|
+
raise HTTPException(
|
|
38
|
+
status_code=500,
|
|
39
|
+
detail=str(exc),
|
|
40
|
+
) from exc
|
|
41
|
+
|
|
42
|
+
return {
|
|
43
|
+
"status": "ok",
|
|
44
|
+
"documents_indexed": len(documents),
|
|
45
|
+
"documents": [
|
|
46
|
+
document.filename
|
|
47
|
+
for document in documents
|
|
48
|
+
],
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
@app.post(
|
|
53
|
+
"/ask",
|
|
54
|
+
response_model=ProofRAGResponse,
|
|
55
|
+
)
|
|
56
|
+
def ask(
|
|
57
|
+
request: ProofRAGRequest,
|
|
58
|
+
) -> ProofRAGResponse:
|
|
59
|
+
try:
|
|
60
|
+
return rag.ask(request)
|
|
61
|
+
except PipelineError as exc:
|
|
62
|
+
raise HTTPException(
|
|
63
|
+
status_code=500,
|
|
64
|
+
detail=str(exc),
|
|
65
|
+
) from exc
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import typer
|
|
4
|
+
from rich.console import Console
|
|
5
|
+
from rich.panel import Panel
|
|
6
|
+
|
|
7
|
+
from .pipeline import PipelineError, ProofRAG
|
|
8
|
+
|
|
9
|
+
app = typer.Typer(
|
|
10
|
+
name="proofrag",
|
|
11
|
+
help="Evidence-backed RAG with verified citations.",
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
console = Console()
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@app.command()
|
|
18
|
+
def ingest() -> None:
|
|
19
|
+
"""Index documents into OKF and Chroma."""
|
|
20
|
+
|
|
21
|
+
rag = ProofRAG()
|
|
22
|
+
|
|
23
|
+
try:
|
|
24
|
+
documents = rag.ingest()
|
|
25
|
+
except PipelineError as exc:
|
|
26
|
+
console.print(
|
|
27
|
+
f"[red]Error:[/red] {exc}"
|
|
28
|
+
)
|
|
29
|
+
raise typer.Exit(code=1)
|
|
30
|
+
|
|
31
|
+
console.print(
|
|
32
|
+
f"[green]Indexed {len(documents)} document(s).[/green]"
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
for document in documents:
|
|
36
|
+
console.print(
|
|
37
|
+
f" • {document.filename}"
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
@app.command()
|
|
42
|
+
def ask(
|
|
43
|
+
question: str = typer.Argument(
|
|
44
|
+
...,
|
|
45
|
+
help="Question to ask over the indexed documents.",
|
|
46
|
+
),
|
|
47
|
+
) -> None:
|
|
48
|
+
"""Ask a question and display verified evidence."""
|
|
49
|
+
|
|
50
|
+
rag = ProofRAG()
|
|
51
|
+
|
|
52
|
+
try:
|
|
53
|
+
response = rag.ask(question)
|
|
54
|
+
except PipelineError as exc:
|
|
55
|
+
console.print(
|
|
56
|
+
f"[red]Error:[/red] {exc}"
|
|
57
|
+
)
|
|
58
|
+
raise typer.Exit(code=1)
|
|
59
|
+
|
|
60
|
+
console.print(
|
|
61
|
+
Panel(
|
|
62
|
+
response.answer,
|
|
63
|
+
title="ProofRAG Answer",
|
|
64
|
+
)
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
console.print(
|
|
68
|
+
f"Verified: {response.verified}"
|
|
69
|
+
)
|
|
70
|
+
console.print(
|
|
71
|
+
f"Refused: {response.refused}"
|
|
72
|
+
)
|
|
73
|
+
console.print(
|
|
74
|
+
f"Evidence: {response.evidence_count}"
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
for citation in response.citations:
|
|
78
|
+
console.print()
|
|
79
|
+
console.print(
|
|
80
|
+
f"[bold]Source:[/bold] {citation.filename}"
|
|
81
|
+
)
|
|
82
|
+
console.print(
|
|
83
|
+
f"[bold]Page:[/bold] {citation.page_number}"
|
|
84
|
+
)
|
|
85
|
+
console.print(
|
|
86
|
+
f"[bold]Status:[/bold] "
|
|
87
|
+
f"{citation.verification_status.value}"
|
|
88
|
+
)
|
|
89
|
+
console.print(
|
|
90
|
+
f"[bold]Evidence:[/bold] "
|
|
91
|
+
f"{citation.snippet}"
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
if __name__ == "__main__":
|
|
96
|
+
app()
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from functools import lru_cache
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
from pydantic import Field
|
|
7
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Settings(BaseSettings):
|
|
14
|
+
"""Application configuration for ProofRAG."""
|
|
15
|
+
|
|
16
|
+
model_config = SettingsConfigDict(
|
|
17
|
+
env_file=PROJECT_ROOT / ".env",
|
|
18
|
+
env_file_encoding="utf-8",
|
|
19
|
+
case_sensitive=False,
|
|
20
|
+
extra="ignore",
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
# ------------------------------------------------------------------
|
|
24
|
+
# Application
|
|
25
|
+
# ------------------------------------------------------------------
|
|
26
|
+
|
|
27
|
+
app_name: str = "ProofRAG"
|
|
28
|
+
app_version: str = "0.1.0"
|
|
29
|
+
environment: str = "development"
|
|
30
|
+
|
|
31
|
+
# ------------------------------------------------------------------
|
|
32
|
+
# Paths
|
|
33
|
+
# ------------------------------------------------------------------
|
|
34
|
+
|
|
35
|
+
data_dir: Path = PROJECT_ROOT / "data"
|
|
36
|
+
documents_dir: Path = PROJECT_ROOT / "data" / "documents"
|
|
37
|
+
okf_dir: Path = PROJECT_ROOT / "data" / "okf"
|
|
38
|
+
chroma_dir: Path = PROJECT_ROOT / "data" / "chroma"
|
|
39
|
+
|
|
40
|
+
# ------------------------------------------------------------------
|
|
41
|
+
# LLM
|
|
42
|
+
# ------------------------------------------------------------------
|
|
43
|
+
|
|
44
|
+
groq_api_key: str = Field(default="", repr=False)
|
|
45
|
+
llm_model: str = "openai/gpt-oss-120b"
|
|
46
|
+
llm_temperature: float = 0.0
|
|
47
|
+
llm_timeout: int = 60
|
|
48
|
+
|
|
49
|
+
# ------------------------------------------------------------------
|
|
50
|
+
# Embeddings
|
|
51
|
+
# ------------------------------------------------------------------
|
|
52
|
+
|
|
53
|
+
embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
|
|
54
|
+
|
|
55
|
+
# ------------------------------------------------------------------
|
|
56
|
+
# Chunking
|
|
57
|
+
# ------------------------------------------------------------------
|
|
58
|
+
|
|
59
|
+
chunk_size: int = 1000
|
|
60
|
+
chunk_overlap: int = 150
|
|
61
|
+
|
|
62
|
+
# ------------------------------------------------------------------
|
|
63
|
+
# Retrieval
|
|
64
|
+
# ------------------------------------------------------------------
|
|
65
|
+
|
|
66
|
+
retrieval_top_k: int = 5
|
|
67
|
+
retrieval_score_threshold: float = 0.35
|
|
68
|
+
|
|
69
|
+
# ------------------------------------------------------------------
|
|
70
|
+
# Citation verification
|
|
71
|
+
# ------------------------------------------------------------------
|
|
72
|
+
|
|
73
|
+
citation_keyword_threshold: float = 0.20
|
|
74
|
+
citation_embedding_threshold: float = 0.60
|
|
75
|
+
|
|
76
|
+
# ------------------------------------------------------------------
|
|
77
|
+
# API
|
|
78
|
+
# ------------------------------------------------------------------
|
|
79
|
+
|
|
80
|
+
api_host: str = "127.0.0.1"
|
|
81
|
+
api_port: int = 8000
|
|
82
|
+
|
|
83
|
+
# ------------------------------------------------------------------
|
|
84
|
+
# Safety / limits
|
|
85
|
+
# ------------------------------------------------------------------
|
|
86
|
+
|
|
87
|
+
max_question_length: int = 2000
|
|
88
|
+
max_context_chunks: int = 8
|
|
89
|
+
|
|
90
|
+
def ensure_directories(self) -> None:
|
|
91
|
+
"""Create required runtime directories if they do not exist."""
|
|
92
|
+
|
|
93
|
+
for directory in (
|
|
94
|
+
self.data_dir,
|
|
95
|
+
self.documents_dir,
|
|
96
|
+
self.okf_dir,
|
|
97
|
+
self.chroma_dir,
|
|
98
|
+
):
|
|
99
|
+
directory.mkdir(parents=True, exist_ok=True)
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
@lru_cache(maxsize=1)
|
|
103
|
+
def get_settings() -> Settings:
|
|
104
|
+
"""Return the cached application settings."""
|
|
105
|
+
|
|
106
|
+
settings = Settings()
|
|
107
|
+
settings.ensure_directories()
|
|
108
|
+
return settings
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from functools import lru_cache
|
|
4
|
+
|
|
5
|
+
from langchain_huggingface import HuggingFaceEmbeddings
|
|
6
|
+
|
|
7
|
+
from .config import get_settings
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@lru_cache(maxsize=1)
|
|
11
|
+
def get_embedding_model() -> HuggingFaceEmbeddings:
|
|
12
|
+
"""
|
|
13
|
+
Create and cache the ProofRAG embedding model.
|
|
14
|
+
|
|
15
|
+
The model runs locally through sentence-transformers.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
settings = get_settings()
|
|
19
|
+
|
|
20
|
+
return HuggingFaceEmbeddings(
|
|
21
|
+
model_name=settings.embedding_model,
|
|
22
|
+
model_kwargs={
|
|
23
|
+
"device": "cpu",
|
|
24
|
+
},
|
|
25
|
+
encode_kwargs={
|
|
26
|
+
"normalize_embeddings": True,
|
|
27
|
+
},
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def embed_text(text: str) -> list[float]:
|
|
32
|
+
"""Generate an embedding for a single text."""
|
|
33
|
+
|
|
34
|
+
if not text.strip():
|
|
35
|
+
raise ValueError("Cannot embed empty text.")
|
|
36
|
+
|
|
37
|
+
return get_embedding_model().embed_query(text)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def embed_documents(
|
|
41
|
+
texts: list[str],
|
|
42
|
+
) -> list[list[float]]:
|
|
43
|
+
"""Generate embeddings for multiple documents."""
|
|
44
|
+
|
|
45
|
+
if not texts:
|
|
46
|
+
return []
|
|
47
|
+
|
|
48
|
+
if any(not text.strip() for text in texts):
|
|
49
|
+
raise ValueError(
|
|
50
|
+
"Document texts cannot contain empty strings."
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
return get_embedding_model().embed_documents(texts)
|