Update models.py for Gemini API correct integration
If you use "ChatGoogleGenerativeAI" instead of "GoogleGenerativeAI", the Agent will output a loop of "AttributeError: 'NoneType' object has no attribute 'stream_generate_content'" }"
Alessandro committed
Sep 7, 2024 at 02:15 UTC
7b005c9f1a7b5f9402bb6c1ca746ff9fb8c95892
1 file changed
+2
-2
models.py
+2
-2
@@ -6,7 +6,7 @@ from langchain_community.embeddings import OllamaEmbeddings
6
from langchain_anthropic import ChatAnthropic
7
from langchain_groq import ChatGroq
8
from langchain_huggingface import HuggingFaceEmbeddings
9
-from langchain_google_genai import ChatGoogleGenerativeAI, HarmBlockThreshold, HarmCategory
9
+from langchain_google_genai import GoogleGenerativeAI, HarmBlockThreshold, HarmCategory
10
from pydantic.v1.types import SecretStr
11
12
@@ -76,7 +76,7 @@ def get_azure_openai_embedding(deployment_name:str, api_key=None, azure_endpoint
76
# Google models
77
def get_google_chat(model_name:str, api_key=None, temperature=DEFAULT_TEMPERATURE):
78
api_key = api_key or get_api_key("google")
79
- return ChatGoogleGenerativeAI(model=model_name, temperature=temperature, google_api_key=api_key, safety_settings={HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE }) # type: ignore
79
+ return GoogleGenerativeAI(model=model_name, temperature=temperature, google_api_key=api_key, safety_settings={HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE }) # type: ignore
80
81
# Groq models
82
def get_groq_chat(model_name:str, api_key=None, temperature=DEFAULT_TEMPERATURE):