Skip to content

Commit

Permalink
fix import ordering
Browse files Browse the repository at this point in the history
  • Loading branch information
rahul2393 committed Feb 19, 2024
1 parent 34430b0 commit 1f9b361
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 5 deletions.
7 changes: 3 additions & 4 deletions src/langchain_google_spanner/chat_message_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,14 @@
"""Cloud Spanner-based chat message history"""
from __future__ import annotations

import json
from typing import List, Optional

from google.cloud import spanner
from google.cloud.spanner_admin_database_v1.types import DatabaseDialect # type: ignore
from google.cloud.spanner_v1 import param_types
from google.cloud.spanner_v1.data_types import JsonObject
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict
from langchain_core.messages import BaseMessage, messages_from_dict

OPERATION_TIMEOUT_SECONDS = 240

Expand Down Expand Up @@ -70,7 +69,7 @@ def _verify_schema(self) -> None:
table with valid schema.
"""
# check table exists
column_names = []
column_names = [] # type: List[str]
with self.database.snapshot() as snapshot:
results = snapshot.execute_sql(
f"SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.columns WHERE table_name = '{self.table_name}'"
Expand Down Expand Up @@ -148,7 +147,7 @@ def messages(self) -> List[BaseMessage]: # type: ignore
params=param,
param_types=param_type,
)
items = []
items = [] # type: List[dict]
for row in results:
items.append({"data": row[0], "type": row[0]["type"]})
messages = messages_from_dict(items)
Expand Down
4 changes: 3 additions & 1 deletion tests/integration/test_spanner_chat_message_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.


import os

import pytest # noqa
from google.cloud.spanner import Client # type: ignore
from langchain_core.messages.ai import AIMessage
Expand All @@ -22,7 +24,7 @@

project_id = os.environ["PROJECT_ID"]
instance_id = os.environ["INSTANCE_ID"]
table_name = os.environ["TABLE_NAME"]
table_name = os.environ["TABLE_NAME"].replace("-", "_")

OPERATION_TIMEOUT_SECONDS = 240

Expand Down

0 comments on commit 1f9b361

Please sign in to comment.