knowledge-grove 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.
Files changed (36) hide show
  1. knowledge_grove-0.1.0/LICENSE +21 -0
  2. knowledge_grove-0.1.0/PKG-INFO +169 -0
  3. knowledge_grove-0.1.0/README.md +128 -0
  4. knowledge_grove-0.1.0/pyproject.toml +35 -0
  5. knowledge_grove-0.1.0/setup.cfg +4 -0
  6. knowledge_grove-0.1.0/src/knowledge_grove/__init__.py +18 -0
  7. knowledge_grove-0.1.0/src/knowledge_grove/cli.py +268 -0
  8. knowledge_grove-0.1.0/src/knowledge_grove/constants.py +34 -0
  9. knowledge_grove-0.1.0/src/knowledge_grove/crud.py +431 -0
  10. knowledge_grove-0.1.0/src/knowledge_grove/db.py +23 -0
  11. knowledge_grove-0.1.0/src/knowledge_grove/mcp_server.py +501 -0
  12. knowledge_grove-0.1.0/src/knowledge_grove/migrations/__init__.py +0 -0
  13. knowledge_grove-0.1.0/src/knowledge_grove/migrations/env.py +85 -0
  14. knowledge_grove-0.1.0/src/knowledge_grove/migrations/versions/initial_schema.py +295 -0
  15. knowledge_grove-0.1.0/src/knowledge_grove/models.py +209 -0
  16. knowledge_grove-0.1.0/src/knowledge_grove/search.py +174 -0
  17. knowledge_grove-0.1.0/src/knowledge_grove/utils/__init__.py +0 -0
  18. knowledge_grove-0.1.0/src/knowledge_grove/utils/chunking.py +128 -0
  19. knowledge_grove-0.1.0/src/knowledge_grove/utils/embedding.py +26 -0
  20. knowledge_grove-0.1.0/src/knowledge_grove/utils/hashing.py +12 -0
  21. knowledge_grove-0.1.0/src/knowledge_grove/utils/input_output.py +78 -0
  22. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/PKG-INFO +169 -0
  23. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/SOURCES.txt +34 -0
  24. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/dependency_links.txt +1 -0
  25. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/entry_points.txt +2 -0
  26. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/requires.txt +11 -0
  27. knowledge_grove-0.1.0/src/knowledge_grove.egg-info/top_level.txt +1 -0
  28. knowledge_grove-0.1.0/tests/test_chunking.py +206 -0
  29. knowledge_grove-0.1.0/tests/test_cli.py +541 -0
  30. knowledge_grove-0.1.0/tests/test_crud.py +1122 -0
  31. knowledge_grove-0.1.0/tests/test_hashing.py +15 -0
  32. knowledge_grove-0.1.0/tests/test_input_output.py +255 -0
  33. knowledge_grove-0.1.0/tests/test_integration.py +141 -0
  34. knowledge_grove-0.1.0/tests/test_mcp_server.py +550 -0
  35. knowledge_grove-0.1.0/tests/test_rls.py +129 -0
  36. knowledge_grove-0.1.0/tests/test_search.py +454 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 konradquam
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,169 @@
1
+ Metadata-Version: 2.4
2
+ Name: knowledge-grove
3
+ Version: 0.1.0
4
+ Summary: A Postgres-backed knowledge schema and SDK for shared agent context
5
+ License: MIT License
6
+
7
+ Copyright (c) 2026 konradquam
8
+
9
+ Permission is hereby granted, free of charge, to any person obtaining a copy
10
+ of this software and associated documentation files (the "Software"), to deal
11
+ in the Software without restriction, including without limitation the rights
12
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13
+ copies of the Software, and to permit persons to whom the Software is
14
+ furnished to do so, subject to the following conditions:
15
+
16
+ The above copyright notice and this permission notice shall be included in all
17
+ copies or substantial portions of the Software.
18
+
19
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
25
+ SOFTWARE.
26
+
27
+ Requires-Python: >=3.11
28
+ Description-Content-Type: text/markdown
29
+ License-File: LICENSE
30
+ Requires-Dist: sqlalchemy>=2.0
31
+ Requires-Dist: alembic>=1.13
32
+ Requires-Dist: psycopg[binary]>=3.1
33
+ Requires-Dist: pgvector>=0.2
34
+ Requires-Dist: sentence_transformers
35
+ Requires-Dist: mcp>=2.0
36
+ Requires-Dist: sqlparse>=0.5
37
+ Provides-Extra: test
38
+ Requires-Dist: pytest>=8.0; extra == "test"
39
+ Requires-Dist: testcontainers[postgres]>=4.0; extra == "test"
40
+ Dynamic: license-file
41
+
42
+ # Knowledge Grove
43
+
44
+ A Postgres-backed knowledge schema and SDK that lets agents — in one repo and across repos — store, discover, and act on shared context: document chunks, embeddings, exact-match tags, and a graph of links between them, including links out to executable tools.
45
+
46
+ **Status:** Core schema, SDK, CLI, and MCP server are implemented and tested (real Postgres, not mocks). A few sections of the original design are still open — see [Implementation status](knowledge-grove-design.md#implementation-status) in the design doc for the precise list.
47
+
48
+ ## Design commitments
49
+
50
+ - **Postgres is the only system of record.** Vector search (`pgvector`), full-text search, exact-match lookups, and the relationship graph all live in one transactional database.
51
+ - **The schema is normalized.** Every relationship — a tag, a link, an access grant — is its own row in its own table.
52
+ - **Access control is enforced by Postgres itself**, via row-level security tied to each agent's own database role. Each agent authenticates with its own role and credentials — never one shared service account.
53
+
54
+ ## Core data model
55
+
56
+ Four tables carry the whole system: `documents` (one row per chunk, with content, a content hash for exact-match dedup, an embedding, and full-text search columns), `document_tags` (exact-match labels), `edges` (a generic relationship graph — next/prev, source, tool links, related, supersedes), and `document_access` (per-document read/write grants). A fifth, `retrieval_feedback`, logs how documents perform for a query, for later ranking-weight tuning.
57
+
58
+ ## Entry points
59
+
60
+ Five ways to land in the graph, each catching a case the others miss: direct ID lookup, tag matching, ILIKE/trigram search, full-text search, and embedding similarity — fused via weighted Reciprocal Rank Fusion. (A bounded, personalized-PageRank walk over the edge graph to expand past those entrance nodes is designed but not yet implemented — see the status link above.)
61
+
62
+ ## Installation
63
+
64
+ Not yet published to PyPI. Install from source:
65
+
66
+ ```bash
67
+ git clone https://github.com/konradquam/knowledge-grove.git
68
+ cd knowledge-grove
69
+ pip install . # or: pip install -e . for an editable install
70
+ ```
71
+
72
+ Requires Python 3.11+ and a Postgres server with the `pgvector` and `pg_trgm` extensions available (the bootstrapping step below creates them for you if your connecting role has privilege to).
73
+
74
+ ## Quickstart
75
+
76
+ **1. Stand up the schema.** Run once per database, using a role with `CREATE EXTENSION`/table-owner privileges — never the role an ordinary agent connects as, since table owners bypass row-level security by default.
77
+
78
+ ```bash
79
+ export KNOWLEDGE_GROVE_DSN="postgresql+psycopg://admin_role:password@host:5432/dbname"
80
+ knowledge-grove init-db
81
+ ```
82
+
83
+ **2. Provision an agent.** Each agent gets its own Postgres role, added to the `shared_reader` group so it can read whatever's been shared into that group by default. Still using the admin DSN:
84
+
85
+ ```bash
86
+ knowledge-grove create-agent-role my_agent
87
+ # Role 'my_agent' created.
88
+ # Agent DSN: postgresql+psycopg://my_agent:<generated-password>@host:5432/dbname
89
+ ```
90
+
91
+ Save that returned DSN somewhere your agent can read it (an env var, a secrets manager — resolving the secret is your project's job, not this package's; see §15 of the design doc).
92
+
93
+ **3. Use it.** Either point an MCP-capable agent at the server (below), ingest existing files from a shell, or call the SDK directly from Python.
94
+
95
+ ```bash
96
+ export KNOWLEDGE_GROVE_DSN="postgresql+psycopg://my_agent:<password>@host:5432/dbname"
97
+ knowledge-grove ingest README.md docs/notes.py
98
+ ```
99
+
100
+ ## CLI reference
101
+
102
+ All commands read their connection string from `KNOWLEDGE_GROVE_DSN`.
103
+
104
+ | Command | Purpose |
105
+ |---|---|
106
+ | `init-db` | Run the bundled Alembic migrations: creates the schema, indexes, and RLS policies. Needs an admin/setup role. |
107
+ | `create-agent-role <name> [--password PW]` | Provision a new agent's Postgres role and `shared_reader` membership. Prompts for a password (hidden) if `--password` is omitted. Needs an admin/setup role. |
108
+ | `ingest <files...> [--source-url URL...] [--content-type {markdown,python,sql}...] [--roles JSON] [-y]` | Chunk and add one or more files as documents. `--source-url`/`--content-type` are given once per file (positionally matched); each defaults respectively to the file's own name and a guess from its extension. Re-ingesting a file under the same `source_url` reconciles against what's already there (a no-op if unchanged, a full replace if not) rather than duplicating it — warns and asks for confirmation first unless `-y` is given. `--roles` is a JSON object (e.g. `'{"shared_reader": ["read"]}'`, the default) applied to every document ingested in the call. |
109
+
110
+ ## MCP server
111
+
112
+ `knowledge_grove.mcp_server` exposes the SDK as MCP tools over stdio (the default transport). Run it directly, or point an MCP-capable client at it:
113
+
114
+ ```json
115
+ {
116
+ "mcpServers": {
117
+ "knowledge-grove": {
118
+ "command": "/path/to/venv/bin/python",
119
+ "args": ["-m", "knowledge_grove.mcp_server"],
120
+ "env": { "KNOWLEDGE_GROVE_DSN": "postgresql+psycopg://my_agent:<password>@host:5432/dbname" }
121
+ }
122
+ }
123
+ }
124
+ ```
125
+
126
+ Each entry in a client's MCP config gets its own subprocess, so each agent identity that needs its own Postgres role should get its own entry with its own DSN — one running server process, one fixed identity for its lifetime; nothing multiplexes several agents through a single connection.
127
+
128
+ Tools exposed:
129
+
130
+ | Tool | Purpose |
131
+ |---|---|
132
+ | `gather_context` | Fused search across all four entry points; the default way to look for existing context. |
133
+ | `add_document` | Add one already-written chunk. |
134
+ | `add_sequential_documents` | Add several chunks in order, auto-linked with `prev` edges. |
135
+ | `add_authored_chunks` | Add several chunks in order, plus arbitrary extra edges (to existing documents, to each other, or to a URL) in the same call. |
136
+ | `get_by_id` | Fetch a document by id. |
137
+ | `get_edges` | Outgoing edges from a document (one hop). |
138
+ | `update_document` | Create a new revision of a document (old one kept, flagged deprecated, linked via `supersedes`). |
139
+ | `add_tag` | Attach an exact-match tag. |
140
+ | `add_edge` | Attach a relationship to another document or an external URL. |
141
+ | `grant_access` / `revoke_access` | Manage a document's access grants. |
142
+ | `log_feedback` | Record how a document performed for a query. |
143
+
144
+ Every tool's full description (usage guidance, argument shapes, when to prefer one over another) is visible to any MCP client that lists tools, and is worth reading directly in `src/knowledge_grove/mcp_server.py` if you're integrating one.
145
+
146
+ ## Using the SDK directly
147
+
148
+ Everything the MCP server exposes is a thin wrapper over `knowledge_grove.crud` and `knowledge_grove.search` — call those directly from Python for non-MCP integrations (e.g. a Temporal activity):
149
+
150
+ ```python
151
+ from knowledge_grove.db import get_engine, get_session
152
+ from knowledge_grove import crud, search
153
+
154
+ engine = get_engine("postgresql+psycopg://my_agent:password@host:5432/dbname")
155
+ session = get_session(engine)
156
+
157
+ doc = crud.add_document(session, content="Retries should use exponential backoff.", owner_agent="my_agent")
158
+ session.commit()
159
+
160
+ hits = search.gather_context(session, query_text="how do retries work", pattern="retry")
161
+ ```
162
+
163
+ ## Design document
164
+
165
+ See [knowledge-grove-design.md](knowledge-grove-design.md) for the full design rationale — auth model, ranking math, chunking strategy, bootstrapping, and the [Implementation status](knowledge-grove-design.md#implementation-status) section tracking what's built against what's still designed-but-open.
166
+
167
+ ## License
168
+
169
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,128 @@
1
+ # Knowledge Grove
2
+
3
+ A Postgres-backed knowledge schema and SDK that lets agents — in one repo and across repos — store, discover, and act on shared context: document chunks, embeddings, exact-match tags, and a graph of links between them, including links out to executable tools.
4
+
5
+ **Status:** Core schema, SDK, CLI, and MCP server are implemented and tested (real Postgres, not mocks). A few sections of the original design are still open — see [Implementation status](knowledge-grove-design.md#implementation-status) in the design doc for the precise list.
6
+
7
+ ## Design commitments
8
+
9
+ - **Postgres is the only system of record.** Vector search (`pgvector`), full-text search, exact-match lookups, and the relationship graph all live in one transactional database.
10
+ - **The schema is normalized.** Every relationship — a tag, a link, an access grant — is its own row in its own table.
11
+ - **Access control is enforced by Postgres itself**, via row-level security tied to each agent's own database role. Each agent authenticates with its own role and credentials — never one shared service account.
12
+
13
+ ## Core data model
14
+
15
+ Four tables carry the whole system: `documents` (one row per chunk, with content, a content hash for exact-match dedup, an embedding, and full-text search columns), `document_tags` (exact-match labels), `edges` (a generic relationship graph — next/prev, source, tool links, related, supersedes), and `document_access` (per-document read/write grants). A fifth, `retrieval_feedback`, logs how documents perform for a query, for later ranking-weight tuning.
16
+
17
+ ## Entry points
18
+
19
+ Five ways to land in the graph, each catching a case the others miss: direct ID lookup, tag matching, ILIKE/trigram search, full-text search, and embedding similarity — fused via weighted Reciprocal Rank Fusion. (A bounded, personalized-PageRank walk over the edge graph to expand past those entrance nodes is designed but not yet implemented — see the status link above.)
20
+
21
+ ## Installation
22
+
23
+ Not yet published to PyPI. Install from source:
24
+
25
+ ```bash
26
+ git clone https://github.com/konradquam/knowledge-grove.git
27
+ cd knowledge-grove
28
+ pip install . # or: pip install -e . for an editable install
29
+ ```
30
+
31
+ Requires Python 3.11+ and a Postgres server with the `pgvector` and `pg_trgm` extensions available (the bootstrapping step below creates them for you if your connecting role has privilege to).
32
+
33
+ ## Quickstart
34
+
35
+ **1. Stand up the schema.** Run once per database, using a role with `CREATE EXTENSION`/table-owner privileges — never the role an ordinary agent connects as, since table owners bypass row-level security by default.
36
+
37
+ ```bash
38
+ export KNOWLEDGE_GROVE_DSN="postgresql+psycopg://admin_role:password@host:5432/dbname"
39
+ knowledge-grove init-db
40
+ ```
41
+
42
+ **2. Provision an agent.** Each agent gets its own Postgres role, added to the `shared_reader` group so it can read whatever's been shared into that group by default. Still using the admin DSN:
43
+
44
+ ```bash
45
+ knowledge-grove create-agent-role my_agent
46
+ # Role 'my_agent' created.
47
+ # Agent DSN: postgresql+psycopg://my_agent:<generated-password>@host:5432/dbname
48
+ ```
49
+
50
+ Save that returned DSN somewhere your agent can read it (an env var, a secrets manager — resolving the secret is your project's job, not this package's; see §15 of the design doc).
51
+
52
+ **3. Use it.** Either point an MCP-capable agent at the server (below), ingest existing files from a shell, or call the SDK directly from Python.
53
+
54
+ ```bash
55
+ export KNOWLEDGE_GROVE_DSN="postgresql+psycopg://my_agent:<password>@host:5432/dbname"
56
+ knowledge-grove ingest README.md docs/notes.py
57
+ ```
58
+
59
+ ## CLI reference
60
+
61
+ All commands read their connection string from `KNOWLEDGE_GROVE_DSN`.
62
+
63
+ | Command | Purpose |
64
+ |---|---|
65
+ | `init-db` | Run the bundled Alembic migrations: creates the schema, indexes, and RLS policies. Needs an admin/setup role. |
66
+ | `create-agent-role <name> [--password PW]` | Provision a new agent's Postgres role and `shared_reader` membership. Prompts for a password (hidden) if `--password` is omitted. Needs an admin/setup role. |
67
+ | `ingest <files...> [--source-url URL...] [--content-type {markdown,python,sql}...] [--roles JSON] [-y]` | Chunk and add one or more files as documents. `--source-url`/`--content-type` are given once per file (positionally matched); each defaults respectively to the file's own name and a guess from its extension. Re-ingesting a file under the same `source_url` reconciles against what's already there (a no-op if unchanged, a full replace if not) rather than duplicating it — warns and asks for confirmation first unless `-y` is given. `--roles` is a JSON object (e.g. `'{"shared_reader": ["read"]}'`, the default) applied to every document ingested in the call. |
68
+
69
+ ## MCP server
70
+
71
+ `knowledge_grove.mcp_server` exposes the SDK as MCP tools over stdio (the default transport). Run it directly, or point an MCP-capable client at it:
72
+
73
+ ```json
74
+ {
75
+ "mcpServers": {
76
+ "knowledge-grove": {
77
+ "command": "/path/to/venv/bin/python",
78
+ "args": ["-m", "knowledge_grove.mcp_server"],
79
+ "env": { "KNOWLEDGE_GROVE_DSN": "postgresql+psycopg://my_agent:<password>@host:5432/dbname" }
80
+ }
81
+ }
82
+ }
83
+ ```
84
+
85
+ Each entry in a client's MCP config gets its own subprocess, so each agent identity that needs its own Postgres role should get its own entry with its own DSN — one running server process, one fixed identity for its lifetime; nothing multiplexes several agents through a single connection.
86
+
87
+ Tools exposed:
88
+
89
+ | Tool | Purpose |
90
+ |---|---|
91
+ | `gather_context` | Fused search across all four entry points; the default way to look for existing context. |
92
+ | `add_document` | Add one already-written chunk. |
93
+ | `add_sequential_documents` | Add several chunks in order, auto-linked with `prev` edges. |
94
+ | `add_authored_chunks` | Add several chunks in order, plus arbitrary extra edges (to existing documents, to each other, or to a URL) in the same call. |
95
+ | `get_by_id` | Fetch a document by id. |
96
+ | `get_edges` | Outgoing edges from a document (one hop). |
97
+ | `update_document` | Create a new revision of a document (old one kept, flagged deprecated, linked via `supersedes`). |
98
+ | `add_tag` | Attach an exact-match tag. |
99
+ | `add_edge` | Attach a relationship to another document or an external URL. |
100
+ | `grant_access` / `revoke_access` | Manage a document's access grants. |
101
+ | `log_feedback` | Record how a document performed for a query. |
102
+
103
+ Every tool's full description (usage guidance, argument shapes, when to prefer one over another) is visible to any MCP client that lists tools, and is worth reading directly in `src/knowledge_grove/mcp_server.py` if you're integrating one.
104
+
105
+ ## Using the SDK directly
106
+
107
+ Everything the MCP server exposes is a thin wrapper over `knowledge_grove.crud` and `knowledge_grove.search` — call those directly from Python for non-MCP integrations (e.g. a Temporal activity):
108
+
109
+ ```python
110
+ from knowledge_grove.db import get_engine, get_session
111
+ from knowledge_grove import crud, search
112
+
113
+ engine = get_engine("postgresql+psycopg://my_agent:password@host:5432/dbname")
114
+ session = get_session(engine)
115
+
116
+ doc = crud.add_document(session, content="Retries should use exponential backoff.", owner_agent="my_agent")
117
+ session.commit()
118
+
119
+ hits = search.gather_context(session, query_text="how do retries work", pattern="retry")
120
+ ```
121
+
122
+ ## Design document
123
+
124
+ See [knowledge-grove-design.md](knowledge-grove-design.md) for the full design rationale — auth model, ranking math, chunking strategy, bootstrapping, and the [Implementation status](knowledge-grove-design.md#implementation-status) section tracking what's built against what's still designed-but-open.
125
+
126
+ ## License
127
+
128
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,35 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "knowledge-grove"
7
+ version = "0.1.0"
8
+ description = "A Postgres-backed knowledge schema and SDK for shared agent context"
9
+ readme = "README.md"
10
+ license = { file = "LICENSE" }
11
+ requires-python = ">=3.11"
12
+ dependencies = [
13
+ "sqlalchemy>=2.0",
14
+ "alembic>=1.13",
15
+ "psycopg[binary]>=3.1",
16
+ "pgvector>=0.2",
17
+ "sentence_transformers",
18
+ "mcp>=2.0",
19
+ "sqlparse>=0.5",
20
+ ]
21
+
22
+ [project.scripts]
23
+ knowledge-grove = "knowledge_grove.cli:main"
24
+
25
+ [project.optional-dependencies]
26
+ test = [
27
+ "pytest>=8.0",
28
+ "testcontainers[postgres]>=4.0",
29
+ ]
30
+
31
+ [tool.setuptools.packages.find]
32
+ where = ["src"]
33
+
34
+ [tool.pytest.ini_options]
35
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,18 @@
1
+ from knowledge_grove.db import get_engine, get_session
2
+ from knowledge_grove.models import (
3
+ Document,
4
+ DocumentTag,
5
+ Edge,
6
+ DocumentAccess,
7
+ RetrievalFeedback,
8
+ )
9
+
10
+ __all__ = [
11
+ "get_engine",
12
+ "get_session",
13
+ "Document",
14
+ "DocumentTag",
15
+ "Edge",
16
+ "DocumentAccess",
17
+ "RetrievalFeedback",
18
+ ]
@@ -0,0 +1,268 @@
1
+ """CLI commands: bootstrapping (`init-db`, `create-agent-role`, §14 of the
2
+ design doc, ops-facing one-time-per-database/per-agent setup) and content
3
+ ingestion (`ingest`, a thin wrapper over add_file_as_document/§13 dedup for
4
+ use from a shell rather than another agent's own code).
5
+ """
6
+ import argparse
7
+ import getpass
8
+ import json
9
+ import os
10
+ from pathlib import Path
11
+
12
+ from alembic import command
13
+ from alembic.config import Config
14
+ from psycopg import sql
15
+ from sqlalchemy import create_engine, func, select, text
16
+ from sqlalchemy.engine import make_url
17
+
18
+ import knowledge_grove
19
+ from knowledge_grove.constants import ContentType, SHARED_READER
20
+ from knowledge_grove.db import get_engine, get_session
21
+ from knowledge_grove.models import Document
22
+ from knowledge_grove.utils.input_output import add_file_as_document, detect_content_type
23
+
24
+
25
+ def _alembic_config(dsn: str) -> Config:
26
+ """Build an Alembic Config pointing at this package's own bundled
27
+ migrations directly, rather than relying on alembic.ini being present on
28
+ disk -- that file lives at the repo root, outside the installed package,
29
+ so it won't exist for a real (non-editable) install of knowledge-grove
30
+ used from another project.
31
+ """
32
+ migrations_dir = Path(knowledge_grove.__file__).parent / "migrations"
33
+ cfg = Config()
34
+ cfg.set_main_option("script_location", str(migrations_dir))
35
+ cfg.set_main_option("sqlalchemy.url", dsn)
36
+ return cfg
37
+
38
+
39
+ def init_db(dsn: str) -> None:
40
+ """Run every bundled migration up to head against `dsn`.
41
+
42
+ Must be run by a role with CREATE EXTENSION / CREATE POLICY / table-owner
43
+ privileges -- never the role an ordinary agent connects as, since table
44
+ owners bypass RLS by default (see the initial migration's own docstring).
45
+ """
46
+ command.upgrade(_alembic_config(dsn), "head")
47
+
48
+
49
+ def create_agent_role(dsn: str, role_name: str, password: str) -> str:
50
+ """Provision a new agent's Postgres role: LOGIN credentials, the base
51
+ grants every agent needs, and membership in `shared_reader` so it can
52
+ read whatever's been shared into that group by default.
53
+
54
+ `dsn` must belong to a role with privileges to create roles and grant
55
+ table access (the same setup-only role `init_db` requires), not an
56
+ ordinary agent role. Returns the DSN the new agent should connect with.
57
+
58
+ Role names and the password can't be passed as ordinary bind parameters
59
+ -- CREATE ROLE / GRANT are utility statements, not DML, so Postgres
60
+ doesn't accept protocol-level placeholders for them. `psycopg.sql`
61
+ composes them safely instead (proper identifier quoting for the role
62
+ name, proper literal escaping for the password).
63
+ """
64
+ engine = create_engine(dsn)
65
+ try:
66
+ with engine.begin() as conn:
67
+ cur = conn.connection.dbapi_connection.cursor()
68
+ cur.execute(
69
+ sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(
70
+ sql.Identifier(role_name), sql.Literal(password)
71
+ )
72
+ )
73
+ cur.execute(
74
+ sql.SQL(
75
+ "GRANT SELECT, INSERT, UPDATE, DELETE ON "
76
+ "documents, document_tags, edges, document_access TO {}"
77
+ ).format(sql.Identifier(role_name))
78
+ )
79
+ cur.execute(
80
+ sql.SQL("GRANT SELECT, INSERT ON retrieval_feedback TO {}").format(
81
+ sql.Identifier(role_name)
82
+ )
83
+ )
84
+ cur.execute(
85
+ sql.SQL("GRANT {} TO {}").format(
86
+ sql.Identifier(SHARED_READER), sql.Identifier(role_name)
87
+ )
88
+ )
89
+ finally:
90
+ engine.dispose()
91
+
92
+ agent_url = make_url(dsn).set(username=role_name, password=password)
93
+ return agent_url.render_as_string(hide_password=False)
94
+
95
+
96
+ def _current_agent(session) -> str:
97
+ """The Postgres role this connection is actually authenticated as -- see
98
+ mcp_server.py's identical helper for why document ownership always comes
99
+ from the connection itself, never a caller-supplied argument.
100
+ """
101
+ return session.execute(text("SELECT current_user")).scalar()
102
+
103
+
104
+ def _active_document_count(session, source_url: str) -> int:
105
+ return session.scalar(
106
+ select(func.count()).select_from(Document).where(
107
+ Document.source_url == source_url, Document.deprecated.is_(False)
108
+ )
109
+ )
110
+
111
+
112
+ def ingest_files(
113
+ dsn: str,
114
+ file_paths: list[str],
115
+ source_urls: list[str | None] | None = None,
116
+ content_types: list[str | None] | None = None,
117
+ roles: dict[str, list[str]] | None = None,
118
+ assume_yes: bool = False,
119
+ ) -> None:
120
+ """Ingest one or more files as documents, chunking each and computing
121
+ embeddings (add_file_as_document). `dsn` should be an ordinary agent's
122
+ connection string, not the admin/setup one `init_db`/`create_agent_role`
123
+ need -- ownership is derived from whichever role `dsn` authenticates as.
124
+
125
+ Each file's `source_url` defaults to its own filename (not a real URL,
126
+ just a stable identifier re-ingesting the same file later will match
127
+ again -- see add_raw_document's §13 reconciliation). If a source_url
128
+ already has existing documents, this asks for confirmation before
129
+ proceeding (unless `assume_yes`): re-using an existing source_url for a
130
+ genuinely different file would make add_raw_document treat it as a
131
+ changed revision and deprecate that unrelated content.
132
+
133
+ Each file's `content_type` (which chunker to use) defaults to a guess
134
+ from its extension (detect_content_type) -- `.py` -> python, `.sql` ->
135
+ sql, everything else -> markdown. Pass `content_types` to override that
136
+ per file, matched by position to `file_paths`.
137
+
138
+ `roles` grants other roles access to every document ingested in this
139
+ call (see add_document); if omitted, defaults to {"shared_reader":
140
+ ["read"]} -- readable by every agent in the shared_reader group. Pass
141
+ {} to keep everything ingested here private to the owner only.
142
+ """
143
+ resolved_source_urls = [
144
+ source_urls[i] if source_urls and source_urls[i] else Path(file_paths[i]).name
145
+ for i in range(len(file_paths))
146
+ ]
147
+ resolved_content_types = [
148
+ content_types[i] if content_types else None
149
+ for i in range(len(file_paths))
150
+ ]
151
+
152
+ engine = get_engine(dsn)
153
+ session = get_session(engine)
154
+ try:
155
+ owner_agent = _current_agent(session)
156
+
157
+ for path, source_url, content_type in zip(file_paths, resolved_source_urls, resolved_content_types):
158
+ existing_count = _active_document_count(session, source_url)
159
+ if existing_count > 0 and not assume_yes:
160
+ print(
161
+ f"Warning: {existing_count} existing document(s) already use "
162
+ f"source_url '{source_url}'. Continuing will treat '{path}' as a "
163
+ f"new revision of that same source: identical content is a "
164
+ f"no-op, but different content will deprecate the existing "
165
+ f"chunks. If '{path}' is not actually a revision of that source, "
166
+ f"answer no and re-run with a different --source-url."
167
+ )
168
+ answer = input("Proceed? [y/N] ").strip().lower()
169
+ if answer != "y":
170
+ print(f"Skipped '{path}'.")
171
+ continue
172
+
173
+ docs = add_file_as_document(
174
+ session, path, owner_agent=owner_agent, source_url=source_url,
175
+ content_type=content_type, roles=roles,
176
+ )
177
+ session.commit()
178
+ used_type = content_type or detect_content_type(path)
179
+ print(f"Ingested '{path}' as {len(docs)} chunk(s) ({used_type}) under source_url '{source_url}'.")
180
+ finally:
181
+ session.close()
182
+ engine.dispose()
183
+
184
+
185
+ def main() -> None:
186
+ parser = argparse.ArgumentParser(prog="knowledge-grove")
187
+ subparsers = parser.add_subparsers(dest="command", required=True)
188
+
189
+ subparsers.add_parser(
190
+ "init-db", help="Run the bundled migrations against KNOWLEDGE_GROVE_DSN."
191
+ )
192
+
193
+ role_parser = subparsers.add_parser(
194
+ "create-agent-role", help="Provision a new agent's Postgres role."
195
+ )
196
+ role_parser.add_argument("role_name")
197
+ role_parser.add_argument(
198
+ "--password", help="If omitted, you'll be prompted (not echoed)."
199
+ )
200
+
201
+ ingest_parser = subparsers.add_parser(
202
+ "ingest", help="Add one or more files as documents, chunking each and computing embeddings."
203
+ )
204
+ ingest_parser.add_argument("files", nargs="+", help="Path(s) to the file(s) to ingest.")
205
+ ingest_parser.add_argument(
206
+ "--source-url", action="append", default=None,
207
+ help=(
208
+ "Source identifier, given once per file in the same order as `files`. "
209
+ "Defaults to each file's own filename (not a real URL -- just a stable "
210
+ "identifier so re-ingesting the same file later is recognized as an "
211
+ "update rather than a new, unrelated document)."
212
+ ),
213
+ )
214
+ ingest_parser.add_argument(
215
+ "--content-type", action="append", default=None, choices=list(ContentType),
216
+ help=(
217
+ "Chunker to use, given once per file in the same order as `files`. "
218
+ "Defaults to a guess from each file's extension (.py -> python, "
219
+ ".sql -> sql, everything else -> markdown)."
220
+ ),
221
+ )
222
+ ingest_parser.add_argument(
223
+ "--roles", default=None,
224
+ help=(
225
+ "JSON object granting other roles access to every document ingested "
226
+ "in this call, e.g. '{\"shared_reader\": [\"read\"]}'. Applies to the "
227
+ "whole call, not per file. Defaults to {\"shared_reader\": [\"read\"]} "
228
+ "if omitted; pass '{}' to keep everything private to the owner."
229
+ ),
230
+ )
231
+ ingest_parser.add_argument(
232
+ "-y", "--yes", action="store_true",
233
+ help="Don't prompt for confirmation when a source_url already has existing documents.",
234
+ )
235
+
236
+ args = parser.parse_args()
237
+
238
+ dsn = os.environ.get("KNOWLEDGE_GROVE_DSN")
239
+ if not dsn:
240
+ parser.error("KNOWLEDGE_GROVE_DSN must be set to a connection string")
241
+
242
+ if args.command == "init-db":
243
+ init_db(dsn)
244
+ print("Migrations applied.")
245
+ elif args.command == "create-agent-role":
246
+ password = args.password or getpass.getpass(f"Password for {args.role_name}: ")
247
+ agent_dsn = create_agent_role(dsn, args.role_name, password)
248
+ print(f"Role '{args.role_name}' created.")
249
+ print(f"Agent DSN: {agent_dsn}")
250
+ elif args.command == "ingest":
251
+ if args.source_url and len(args.source_url) != len(args.files):
252
+ parser.error("--source-url must be given once per file, or omitted entirely")
253
+ if args.content_type and len(args.content_type) != len(args.files):
254
+ parser.error("--content-type must be given once per file, or omitted entirely")
255
+ roles = None
256
+ if args.roles is not None:
257
+ try:
258
+ roles = json.loads(args.roles)
259
+ except json.JSONDecodeError as e:
260
+ parser.error(f"--roles must be valid JSON: {e}")
261
+ ingest_files(
262
+ dsn, args.files, source_urls=args.source_url,
263
+ content_types=args.content_type, roles=roles, assume_yes=args.yes,
264
+ )
265
+
266
+
267
+ if __name__ == "__main__":
268
+ main()