Skip to content

Commit 62a02b6

Browse files
committed
chore: add more unit tests
1 parent 76d77a1 commit 62a02b6

File tree

2 files changed

+22
-1
lines changed

2 files changed

+22
-1
lines changed

models/ssm.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import os
1111
from typing import List # ClassVar
1212

13+
# pinecone integration
1314
import pinecone
1415
from langchain.cache import InMemoryCache
1516

@@ -76,7 +77,8 @@ def cached_chat_request(self, system_message: str, human_message: str) -> System
7677
HumanMessage(content=human_message),
7778
]
7879
# pylint: disable=not-callable
79-
return self.chat(messages)
80+
retval = self.chat(messages).content
81+
return retval
8082

8183
def prompt_with_template(self, prompt: PromptTemplate, concept: str, model: str = DEFAULT_MODEL_NAME) -> str:
8284
"""Prompt with template."""

models/tests/test_ssm.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@
44
"""
55
Test integrity of base class.
66
"""
7+
import pinecone
78
import pytest # pylint: disable=unused-import
89
from langchain.chat_models import ChatOpenAI
910
from langchain.embeddings import OpenAIEmbeddings
1011
from langchain.text_splitter import RecursiveCharacterTextSplitter
1112
from langchain.vectorstores.pinecone import Pinecone
1213

14+
from ..const import Credentials
1315
from ..ssm import SalesSupportModel
1416

1517

@@ -29,3 +31,20 @@ def test_02_class_aatribute_types(self):
2931
assert isinstance(ssm.pinecone_index, Pinecone)
3032
assert isinstance(ssm.text_splitter, RecursiveCharacterTextSplitter)
3133
assert isinstance(ssm.openai_embedding, OpenAIEmbeddings)
34+
35+
def test_03_test_openai_connectivity(self):
36+
"""Ensure that we have connectivity to OpenAI."""
37+
38+
ssm = SalesSupportModel()
39+
retval = ssm.cached_chat_request(
40+
"your are a helpful assistant", "please return the value 'CORRECT' in all upper case."
41+
)
42+
assert retval == "CORRECT"
43+
44+
def test_04_test_pinecone_connectivity(self):
45+
"""Ensure that we have connectivity to Pinecone."""
46+
# pylint: disable=broad-except
47+
try:
48+
pinecone.init(api_key=Credentials.PINECONE_API_KEY, environment=Credentials.PINECONE_ENVIRONMENT)
49+
except Exception as e:
50+
assert False, f"pinecone.init() failed with exception: {e}"

0 commit comments

Comments
 (0)