Skip to content

Commit

Permalink
set private attributes with PrivateAttr instead of config
Browse files Browse the repository at this point in the history
  • Loading branch information
ccurme committed Sep 9, 2024
1 parent d9f0342 commit 7d8c6f6
Showing 1 changed file with 9 additions and 11 deletions.
20 changes: 9 additions & 11 deletions libs/community/langchain_google_community/vertex_ai_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
from langchain_core.retrievers import BaseRetriever
from langchain_core.tools import BaseTool
from langchain_core.utils import get_from_dict_or_env
from pydantic import ConfigDict, Field, model_validator
from pydantic import ConfigDict, Field, PrivateAttr, model_validator

from langchain_google_community._utils import get_client_info

Expand Down Expand Up @@ -256,13 +256,12 @@ class VertexAISearchRetriever(BaseRetriever, _BaseVertexAISearchRetriever):
https://cloud.google.com/generative-ai-app-builder/docs/reference/rest/v1beta/BoostSpec
"""

_client: SearchServiceClient
_serving_config: str
_client: SearchServiceClient = PrivateAttr()
_serving_config: str = PrivateAttr()

model_config = ConfigDict(
extra="forbid",
arbitrary_types_allowed=True,
underscore_attrs_are_private=True,
)

def __init__(self, **kwargs: Any) -> None:
Expand Down Expand Up @@ -418,13 +417,12 @@ class VertexAIMultiTurnSearchRetriever(BaseRetriever, _BaseVertexAISearchRetriev
conversation_id: str = "-"
"""Vertex AI Search Conversation ID."""

_client: ConversationalSearchServiceClient
_serving_config: str
_client: ConversationalSearchServiceClient = PrivateAttr()
_serving_config: str = PrivateAttr()

model_config = ConfigDict(
extra="ignore",
arbitrary_types_allowed=True,
underscore_attrs_are_private=True,
)

def __init__(self, **kwargs: Any):
Expand Down Expand Up @@ -501,10 +499,10 @@ class VertexAISearchSummaryTool(BaseTool, VertexAISearchRetriever):
summary_spec_kwargs: Dict[str, Any] = Field(default_factory=dict)
""" Additional kwargs for `SearchRequest.ContentSearchSpec.SummarySpec`"""

class Config(VertexAISearchRetriever.Config):
"""Redefinition to specify that inherits config from `VertexAISearchRetriever`
not BaseTool
"""
model_config = ConfigDict(
extra="forbid",
arbitrary_types_allowed=True,
)

def _get_content_spec_kwargs(self) -> Optional[Dict[str, Any]]:
"""Adds additional summary_spec parameters to the configuration of the search.
Expand Down

0 comments on commit 7d8c6f6

Please sign in to comment.