Skip to content

Update haiway with Postgres proxy #417

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KaQuMiQ merged 1 commit into from
Sep 17, 2025
Merged

Update haiway with Postgres proxy #417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ build-backend = "uv_build"
[project]
name = "draive"
description = "Framework designed to simplify and accelerate the development of LLM-based applications."
version = "0.84.7"
version = "0.85.0"
readme = "README.md"
maintainers = [
{ name = "Kacper Kaliński", email = "kacper.kalinski@miquido.com" },
Expand All @@ -21,7 +21,7 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Application Frameworks",
]
license = { file = "LICENSE" }
dependencies = ["numpy~=2.2", "haiway~=0.31.4"]
dependencies = ["numpy~=2.2", "haiway~=0.32.0"]

[project.urls]
Homepage = "https://miquido.com"
Expand All @@ -47,6 +47,7 @@ opentelemetry = [
"opentelemetry-exporter-otlp-proto-grpc",
]
httpx = ["haiway[httpx]", "httpx"]
postgres = ["haiway[postgres]", "asyncpg"]
Copy link

@coderabbitai coderabbitai bot Sep 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

🧩 Analysis chain

Optional group “postgres”: confirm duplication and pin asyncpg.
If haiway[postgres] already depends on asyncpg, listing asyncpg here duplicates it. If not, please pin a compatible range (e.g., asyncpg>=0.30,<1) to avoid resolver drift. Latest stable of asyncpg is 0.30.0. (pypi.org)

Option A (if haiway extra includes driver):

-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]"]

Option B (if not included):

-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"]

🌐 Web query:

Does the haiway package define a "postgres" extra that installs asyncpg?

💡 Result:

No — the haiway package does not expose a "postgres" extra that installs asyncpg. (haiway's PyPI project page and repository show no extras for a postgres/asyncpg extra). [1][2]

Sources:
[1] haiway on PyPI. [2] miquido/haiway repository (GitHub).

Pin asyncpg in pyproject.toml — haiway has no postgres extra

haiway does not expose a "postgres" extra that installs asyncpg — pin asyncpg to avoid resolver drift.

-postgres = ["haiway[postgres]", "asyncpg"]
+postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"]
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
postgres = ["haiway[postgres]", "asyncpg"]
postgres = ["haiway[postgres]", "asyncpg>=0.30,<1"]
🤖 Prompt for AI Agents 
In pyproject.toml at line 50, the dependency list uses "asyncpg" unpinned while
haiway does not provide a "postgres" extra; update that entry to pin asyncpg to
a specific, compatible version range (for example a maintained minor range like
>=0.26.0,<0.28.0) so the resolver won’t drift — edit the postgres line to
replace the bare "asyncpg" with the chosen version specifier and run poetry
lock/resolve to verify.

dev = [
"bandit~=1.8",
"pyright~=1.1",
Expand Down
2 changes: 2 additions & 0 deletions src/draive/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
HTTPQueryParams,
HTTPResponse,
HTTPStatusCode,
Immutable,
LoggerObservability,
Meta,
MetaTags,
Expand Down Expand Up @@ -260,6 +261,7 @@
"HTTPStatusCode",
"ImageEmbedding",
"ImageGeneration",
"Immutable",
"Instructions",
"InstructionsArgumentDeclaration",
"InstructionsDeclaration",
Expand Down
59 changes: 57 additions & 2 deletions src/draive/gemini/generating.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,10 @@
from haiway import META_EMPTY, MISSING, ObservabilityLevel, as_dict, as_list, ctx

from draive.gemini.api import GeminiAPI
from draive.gemini.config import GeminiConfig, GeminiSafetyConfig
from draive.gemini.config import (
GeminiConfig,
GeminiSafetyConfig,
)
from draive.gemini.utils import unwrap_missing
from draive.models import (
GenerativeModel,
Expand Down Expand Up @@ -268,7 +271,7 @@ async def _completion(
},
)

async def _completion_stream(
async def _completion_stream( # noqa: C901
self,
*,
instructions: ModelInstructions,
Expand Down Expand Up @@ -320,6 +323,10 @@ async def _completion_stream(
contents=request_content,
)

collected_blocks: list[ModelOutputBlock] = []
finish_reason: FinishReason | None = None
finish_message: str | None = None

async for chunk in response_stream:
_record_usage_metrics(
chunk.usage_metadata,
Expand All @@ -330,14 +337,62 @@ async def _completion_stream(
continue

chunk_candidate: Candidate = chunk.candidates[0] # we always request only one
finish_reason = chunk_candidate.finish_reason
finish_message = chunk_candidate.finish_message

if chunk_candidate.safety_ratings:
ctx.record(
ObservabilityLevel.INFO,
event="model.safety.results",
attributes={
"results": [
f"{rating.category} |blocked: {rating.blocked}"
f" |probability:{rating.probability_score}"
f" |severity:{rating.severity_score}"
for rating in chunk_candidate.safety_ratings
if rating.category
],
},
)

if not chunk_candidate.content or not chunk_candidate.content.parts:
continue

for part in chunk_candidate.content.parts:
collected_blocks.extend(_part_as_output_blocks(part))
for element in _part_as_stream_elements(part):
# element is either a MultimodalContentElement or ModelToolRequest
yield element

if finish_reason == FinishReason.SAFETY:
raise ModelOutputFailed(
provider="gemini",
model=config.model,
reason=f"Safety filtering: {finish_message or ''}",
)

if finish_reason == FinishReason.MAX_TOKENS:
raise ModelOutputLimit(
provider="gemini",
model=config.model,
max_output_tokens=unwrap_missing(
config.max_output_tokens,
default=0,
),
content=tuple(collected_blocks),
)

if finish_reason not in (None, FinishReason.STOP):
raise ModelOutputFailed(
provider="gemini",
model=config.model,
reason=(
f"Completion error: {finish_message}"
if finish_message
else "Completion error"
),
)

Comment on lines +367 to +395
Copy link

@coderabbitai coderabbitai bot Sep 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Don’t swallow ModelOutputLimit/ModelOutputFailed raised post-stream.
The broad except Exception will catch the deliberate raises after the stream ends, converting them into ModelOutputFailed and losing semantics. Let our control‑flow exceptions pass through.

Apply a focused guard before the broad catch:

             except ResourceExhausted as exc:
                 ctx.record(
                     ObservabilityLevel.WARNING,
                     event="model.rate_limit",
                     attributes={
                         "model.provider": "gemini",
                         "model.name": config.model,
                     },
                 )
                 # Propagate as ModelRateLimit with randomized backoff window
                 raise ModelRateLimit(
                     provider="gemini",
                     model=config.model,
                     retry_after=random.uniform(*RATE_LIMIT_RETRY_RANGE),  # nosec: B311
                 ) from exc

+            except (ModelOutputFailed, ModelOutputLimit):
+                # Preserve intended control-flow exceptions from end-of-stream checks
+                raise
+
             except Exception as exc:
                 # Convert to ModelOutputFailed for consistency
                 raise ModelOutputFailed(
                     provider="gemini",
                     model=config.model,
                     reason=str(exc),
                 ) from exc

Also applies to: 412-419

🤖 Prompt for AI Agents 
In src/draive/gemini/generating.py around lines 367-395 (and similarly 412-419),
the current broad except Exception block swallows deliberate
ModelOutputLimit/ModelOutputFailed exceptions raised after the stream; update
the exception handler so that it re-raises those control-flow exceptions instead
of converting them. Concretely, change the generic except to capture the
exception as e and add an immediate guard like: if isinstance(e,
(ModelOutputLimit, ModelOutputFailed)): raise; otherwise handle/log the
exception as before; ensure the names are in scope or imported as needed.

except ResourceExhausted as exc:
ctx.record(
ObservabilityLevel.WARNING,
Expand Down
17 changes: 17 additions & 0 deletions src/draive/postgres/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from haiway.postgres import (
Postgres,
PostgresConnection,
PostgresConnectionPool,
PostgresException,
PostgresRow,
PostgresValue,
)

__all__ = (
"Postgres",
"PostgresConnection",
"PostgresConnectionPool",
"PostgresException",
"PostgresRow",
"PostgresValue",
)
Loading