feat(postgis): add PostGIS documentation support - #59
Conversation
|
|
b716454 to
e5d7b77
Compare
murrayju
left a comment
murrayju
left a comment
There was a problem hiding this comment.
Thanks for the submission! A few minor requests, if you don't mind fixing them.
| conn.execute("DROP TABLE IF EXISTS docs.postgis_chunks_tmp CASCADE") | ||
| conn.execute("DROP TABLE IF EXISTS docs.postgis_pages_tmp CASCADE") | ||
|
|
||
| # 创建页面表 |
There was a problem hiding this comment.
Can we use English for all comments, please?
| ) | ||
| """) | ||
|
|
||
| # 创建块表 |
|
|
||
| args = parser.parse_args() | ||
|
|
||
| # 验证数据库存储需求 |
| , sub_chunk_index INTEGER NOT NULL DEFAULT 0 | ||
| , content TEXT NOT NULL | ||
| , metadata JSONB | ||
| , embedding vector(1536) |
There was a problem hiding this comment.
The schema here (and in the existing pg and tiger schemas) hard codes the vector size to 1536. I don't think adding EMBEDDING_DIMENSIONS as an environment variable adds value, unless the schema would be updated/migrated to match. Otherwise, inserts will just fail.
I'd be inclined to keep this fixed as 1536 for now.
| pg_database = os.environ.get("PGDATABASE") | ||
|
|
||
| if all([pg_user, pg_password, pg_host, pg_port, pg_database]): | ||
| return f"postgresql://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_database}" |
There was a problem hiding this comment.
The password is not being url-encoded, so a password containing e.g. @ would break the connection string. I can see this is also an issue in the existing code, so we could file a ticket to address this later if you'd prefer.
| """ | ||
|
|
||
| import argparse | ||
| from dataclasses import dataclass, field |
| from psycopg.sql import SQL, Identifier | ||
| import re | ||
| import requests | ||
| from urllib.parse import urljoin, urlparse |
| import time | ||
|
|
||
| THIS_DIR = Path(__file__).parent.resolve() | ||
| load_dotenv(dotenv_path=os.path.join(THIS_DIR, "..", ".env")) |
There was a problem hiding this comment.
| load_dotenv(dotenv_path=os.path.join(THIS_DIR, "..", ".env")) | |
| load_dotenv(dotenv_path=THIS_DIR.parent / ".env") |
|
Hi @murrayju, thank you for the thorough review! I've addressed all the feedback in my latest commits:
Regarding I've also added a new Thanks again for the review! |
@cbc3929 I understood this to mean that you would remove the Could you also resolve the merge conflict? We switched to biome, and you just need to run |
|
Hi @murrayju, thanks for the follow-up! I've addressed your feedback in my latest commits:
I also noticed and fixed a related issue:
Ready for another look! 🙏 |
|
@cbc3929 sorry for the delay here, I was at an offsite last week. I just merged #68, which unfortunately creates some merge conflicts here. However, it addresses a concern with tool bloat as we add more docs sources to this project. You should be able to fold your additions into the singular I also added migrations for the bm25 indexes, which were missing here as well. Let me know if any of that is unclear! |
Add comprehensive PostGIS documentation scraping and semantic search capabilities: New features: - PostGIS manual scraper (postgis_docs.py) for DocBook HTML documentation - Semantic search API (semantic_search_postgis_docs) - Keyword search API (keyword_search_postgis_docs) - Database migration for postgis_pages and postgis_chunks tables - HNSW vector index for fast similarity search Enhanced embedding configuration: - Add OPENAI_BASE_URL for custom OpenAI-compatible endpoints - Add EMBEDDING_MODEL for custom embedding model selection - Add EMBEDDING_DIMENSIONS for configurable vector dimensions - Support third-party embedding services (e.g., SiliconFlow, Ollama) Updated files: - README.md: Add PostGIS to supported extensions - .env.sample: Document new embedding configuration options - tiger_docs.py, postgres_docs.py: Add custom embedding support
- Replace Chinese comments with English - Remove unused imports (field, urlparse) - Use pathlib style for dotenv path - Fix EMBEDDING_DIMENSIONS to 1536 to match database schema - Add URL encoding for password in connection string
Comprehensive reference covering: - Geometry vs Geography selection guide - Coordinate systems (SRID) best practices - Spatial indexing (GiST, BRIN, SP-GiST) - Table design examples (POI, parcels, GPS tracking) - Performance optimization patterns - Data validation techniques
- Add URL encoding for database passwords in tiger_docs.py and postgres_docs.py to handle special characters like '@' in connection strings - Update .env.sample to clarify that EMBEDDING_DIMENSIONS is fixed at 1536 (removed configurable env var since it must match database schema)
36e8e3b to
1694507
Compare
|
Hi @murrayju, Thank you for your patience and the detailed review feedback. Apologies for not keeping up with the repository — I wasn't aware of the major refactoring in PR #68 until now. I've now rebased onto Changes made:
This should address the tool bloat concern and align with the patterns established in PR #68. Please let me know if there's anything else that needs to be adjusted. Thanks again for your guidance! |
| await client.query(/* sql */ ` | ||
| CREATE INDEX CONCURRENTLY IF NOT EXISTS postgis_chunks_content_idx | ||
| ON ${schema}.postgis_chunks | ||
| USING bm25(content) WITH (text_config='english'); | ||
| `); |
There was a problem hiding this comment.
These should be added to a new migration, not edit an existing migration.
| OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | ||
| OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL") # Optional: custom API endpoint | ||
| EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small") # Default model | ||
| EMBEDDING_DIMENSIONS = int(os.getenv("EMBEDDING_DIMENSIONS", "1536")) # Default dimensions |
There was a problem hiding this comment.
I'd still prefer to declare this as a const, not read the env var (the schema is inflexible)
| OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') | ||
| OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL') # Optional: custom API endpoint | ||
| EMBEDDING_MODEL = os.getenv('EMBEDDING_MODEL', 'text-embedding-3-small') # Default model | ||
| EMBEDDING_DIMENSIONS = int(os.getenv('EMBEDDING_DIMENSIONS', '1536')) # Default dimensions |
- Create separate migration for PostGIS BM25 index instead of editing existing migration - Change EMBEDDING_DIMENSIONS to const in postgres_docs.py and tiger_docs.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
|
Hi @murrayju, Thanks for the feedback! I've addressed all three points:
Let me know if there's anything else! |
| POSTGIS_DOMAIN = "postgis.net" | ||
|
|
||
| # Token counting using tiktoken | ||
| ENC = tiktoken.get_encoding("cl100k_base") |
There was a problem hiding this comment.
| ENC = tiktoken.get_encoding("cl100k_base") | |
| ENC = tiktoken.encoding_for_model(EMBEDDING_MODEL) |
This simplifies code a bit and will prevent issues with tokenizer not matching a non-default model via envvar
| return None | ||
|
|
||
| try: | ||
| time.sleep(self.delay) |
There was a problem hiding this comment.
Is a sleep needed before making a request that has a 30 second timeout?
|
|
||
| # Remove images with data: URLs | ||
| for img in soup.find_all("img", src=True): | ||
| if img["src"].startswith("data:"): |
There was a problem hiding this comment.
may be a good idea to also look at the srcset array?
| def chunk_markdown(self, markdown: str, page: Page) -> list[Chunk]: | ||
| """Split Markdown into chunks based on headers.""" | ||
| chunks = [] | ||
| header_pattern = re.compile(r"^(#{1,3}) (.+)$", re.MULTILINE) |
There was a problem hiding this comment.
given that we are splitting on newlines, we shouldn't need multiline. Also, is there a particular reason to limit the definition of a header to just # - ###?
| header_pattern = re.compile(r"^(#{1,3}) (.+)$", re.MULTILINE) | |
| header_pattern = re.compile(r"^(#+) (.+)$") |
| if current_chunk_lines: | ||
| content = "\n".join(current_chunk_lines).strip() | ||
| if content: | ||
| chunks.append(Chunk( |
There was a problem hiding this comment.
Rather than creating chunks, appending them, then iterating over all of the chunks to then conditionally split them into subchunks, I think it would be more straighforward to create a method like
def create_chunks(idx: int, header: str, header_path: str, content: str) -> list[Chunk]:
chunks = []
tokens = ENC.encode(content)
number_of_chunks = math.ceil(len(tokens) / MAX_CHUNK_TOKENS)
for sub_idx in range(number_of_chunks):
chunk_tokens = tokens[sub_idx * MAX_CHUNK_TOKENS:(sub_idx + 1)*MAX_CHUNK_TOKENS]
chunk_content = ENC.decode(chunk_tokens)
chunks.append(Chunk(
idx=idx,
header=header,
header_path=header_path,
content=chunk_content,
token_count=len(chunk_tokens),
subindex=sub_idx,
))then this code would look like
for line in lines:
match = header_pattern.match(line)
if match:
# Save current chunk
if current_chunk_lines:
content = "\n".join(current_chunk_lines).strip()
if content:
chunks.extend(create_chunks(idx=idx, header=current_header, header_path=header_path.copy(), content=content)
|
|
||
| # Token counting using tiktoken | ||
| ENC = tiktoken.get_encoding("cl100k_base") | ||
| MAX_CHUNK_TOKENS = 7000 |
There was a problem hiding this comment.
openai allows for 8191 tokens per embedding (doc), is there a reason you chose 7000?
| for chunk in chunks: | ||
| # Generate embedding using configurable model and dimensions | ||
| try: | ||
| embedding = ( |
There was a problem hiding this comment.
Structure-wise, I would rather that the save_to_database method was more limited to just the persistence of the data models and not responsible for the logic to create embeddings.
Logic wise, there are a few issues with this:
- there is a limit to the number of tokens that can be embedded in one "document"
- higher chance of hitting rate limits on creating embeddings if you are calling that endpoint separately for every single chunk, rather than combining multiple chunks into a single request.
In order to minimize the risk of rate limiting (and reducing cost), we should combine multiple chunks into a single request
What if a function like this were added
MAX_TOKEN_PER_REQUEST = 300_000
async def get_chunk_embeddings(chunks: list[Chunk | None]) -> list[float]:
# i would just move the client creation to the top of the script
current_batch_token_count = 0
current_batch: list[Chunk] = []
embeddings: float[] = []
chunks.append(None) #signal that we are out of chunks
for chunk in chunks:
add_chunk_to_batch = chunk.token_count + current_batch_token_count < MAX_TOKEN_PER_REQUEST if chunk else False
if add_chunk_to_batch:
current_batch.append(chunk)
current_batch_token_count += chunk.token_count
if not add_chunk_to_batch:
batch_embeddings = client.embeddings.create(input=current_batch)
embeddings.extend(batch_embeddedings)
current_batch = [chunk] if chunk else []
current_batch_token_count = chunk.token_count if chunk else 0
Then you would have a signature for this method like:
def save_to_database(
self,
conn: psycopg.Connection,
page: Page,
chunks: list[Chunk],
embeddings: list[float]
|
|
||
| conn.execute( | ||
| """ | ||
| INSERT INTO docs.postgis_chunks_tmp |
There was a problem hiding this comment.
rather than an execute insert for each chunk, you could pass an array as the parameter into the execute and use unnest in the sql to iterate over each json row.
| ".sr-only", | ||
| ".code-block-copy-button" | ||
| ] | ||
| # Add more domains as needed |
There was a problem hiding this comment.
should probably keep this comment below your postgis.net entry
|
@cbc3929 I am going to merge your changes and will follow up with my own PR with my suggestions. Thank you for your contribution! We apologize for the duration of this review. |
Signed-off-by: Greg Saab <greg@onethirty.one>
## What Refactor `postgres_docs.py` and `postgis_docs.py` so we can share functionality between the two. Other changes: * add vscode launch/tasks for debugging * add ruff rules * add workflows for prod/dev ## Why This is a follow up to a external PR submission -- decision was to merge that PR then address refactoring. ## Testing I tested postgres and postgis locally and verified that the outputs look good. I scanned through all of the chunks and spot checked a bunch of them. I have attached the outputs from both [postgres.json](https://github.com/user-attachments/files/25580373/postgres.json) [postgis.json](https://github.com/user-attachments/files/25580374/postgis.json) ## Related PR #59
4 participants
Summary
This PR adds comprehensive PostGIS documentation support to pg-aiguide, enabling AI coding assistants to provide better guidance for spatial database operations.
New Features
ingest/postgis_docs.py)semantic_search_postgis_docs- Vector similarity search for PostGIS documentationkeyword_search_postgis_docs- BM25 keyword search for PostGIS documentationpostgis_pagesandpostgis_chunkstablesEnhanced Embedding Configuration
Added support for custom embedding providers (beyond OpenAI):
OPENAI_BASE_URL- Custom OpenAI-compatible API endpoint (e.g., Ollama, SiliconFlow)EMBEDDING_MODEL- Custom embedding model nameEMBEDDING_DIMENSIONS- Configurable vector dimensionsThis allows users to use alternative embedding services while maintaining compatibility with the existing database schema.
Testing
Usage