GML-1965 Gemini support and GML-1966 log LLM usage

README.md

@@ -124,7 +124,37 @@ In addition to the `OPENAI_API_KEY`, `llm_model` and `model_name` can be edited
 }
 ```
 
-* GCP
+* Google GenAI
+
+Get your Gemini API key via https://aistudio.google.com/app/apikey.
+
+```json
+{
+    "llm_config": {
+        "embedding_service": {
+            "embedding_model_service": "genai",
+            "model_name": "models/gemini-embedding-exp-03-07",
+            "dimensions": 1536,
+            "authentication_configuration": {
+                "GOOGLE_API_KEY": "YOUR_GOOGLE_API_KEY_HERE"
+            }
+        },
+        "completion_service": {
+            "llm_service": "genai",
+            "llm_model": "gemini-2.5-flash",
+            "authentication_configuration": {
+                "GOOGLE_API_KEY": "YOUR_GOOGLE_API_KEY_HERE"
+            },
+            "model_kwargs": {
+                "temperature": 0
+            },
+            "prompt_path": "./common/prompts/google_gemini/"
+        }
+    }
+}
+```
+
+* GCP VertexAI
 
 Follow the GCP authentication information found here: https://cloud.google.com/docs/authentication/application-default-credentials#GAC and create a Service Account with VertexAI credentials. Then add the following to the docker run command:
 

common/config.py

@@ -9,13 +9,15 @@
     AzureOpenAI_Ada002,
     OpenAI_Embedding,
     VertexAI_PaLM_Embedding,
+    GenAI_Embedding,
 )
 from common.embeddings.tigergraph_embedding_store import TigerGraphEmbeddingStore
 from common.llm_services import (
     AWS_SageMaker_Endpoint,
     AWSBedrock,
     AzureOpenAI,
     GoogleVertexAI,
+    GoogleGenAI,
     Groq,
     HuggingFaceEndpoint,
     LLM_Model,
@@ -62,12 +64,20 @@
 graphrag_config = server_config.get("graphrag_config")
 
 if db_config is None:
-    raise Exception("graphrag_config is not found in SERVER_CONFIG")
+    raise Exception("db_config is not found in SERVER_CONFIG")
 if llm_config is None:
-    raise Exception("graphrag_config is not found in SERVER_CONFIG")
+    raise Exception("llm_config is not found in SERVER_CONFIG")
+
+completion_config = llm_config.get("completion_service")
+if completion_config is None:
+    raise Exception("completion_service is not found in llm_config")
+embedding_config = llm_config.get("embedding_service")
+if embedding_config is None:
+    raise Exception("embedding_service is not found in llm_config")
+embedding_dimension = embedding_config.get("dimensions", 1536)
 
 if graphrag_config is None:
-    graphrag_config = {"reuse_embedding", true}
+    graphrag_config = {"reuse_embedding", True}
 if "chunker" not in graphrag_config:
     graphrag_config["chunker"] = "semantic"
 if "extractor" not in graphrag_config:
@@ -85,6 +95,8 @@
     embedding_service = AzureOpenAI_Ada002(llm_config["embedding_service"])
 elif llm_config["embedding_service"]["embedding_model_service"].lower() == "vertexai":
     embedding_service = VertexAI_PaLM_Embedding(llm_config["embedding_service"])
+elif llm_config["embedding_service"]["embedding_model_service"].lower() == "genai":
+    embedding_service = GenAI_Embedding(llm_config["embedding_service"])
 elif llm_config["embedding_service"]["embedding_model_service"].lower() == "bedrock":
     embedding_service = AWS_Bedrock_Embedding(llm_config["embedding_service"])
 else:
@@ -99,6 +111,8 @@ def get_llm_service(llm_config) -> LLM_Model:
         return AWS_SageMaker_Endpoint(llm_config["completion_service"])
     elif llm_config["completion_service"]["llm_service"].lower() == "vertexai":
         return GoogleVertexAI(llm_config["completion_service"])
+    elif llm_config["completion_service"]["llm_service"].lower() == "genai":
+        return GoogleGenAI(llm_config["completion_service"])
     elif llm_config["completion_service"]["llm_service"].lower() == "bedrock":
         return AWSBedrock(llm_config["completion_service"])
     elif llm_config["completion_service"]["llm_service"].lower() == "groq":
@@ -129,4 +143,3 @@ def get_llm_service(llm_config) -> LLM_Model:
         support_ai_instance=True,
     )
     service_status["embedding_store"] = {"status": "ok", "error": None}
-

common/embeddings/embedding_services.py

@@ -4,6 +4,8 @@
 from typing import List
 
 from langchain.schema.embeddings import Embeddings
+from langchain_community.embeddings.openai import OpenAIEmbeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
 
 from common.logs.log import req_id_cv
 from common.logs.logwriter import LogWriter
@@ -17,7 +19,7 @@ class EmbeddingModel(Embeddings):
     Implements connections to the desired embedding API.
     """
 
-    def __init__(self, config: dict, model_name: str, base_url: str = None):
+    def __init__(self, config: dict, model_name: str):
         """Initialize an EmbeddingModel
         Read JSON config file and export the details as environment variables.
         """
@@ -27,9 +29,8 @@ def __init__(self, config: dict, model_name: str, base_url: str = None):
             ]
         self.embeddings = None
         self.model_name = model_name
-        self.base_url = config.get("base_url")
         LogWriter.info(
-            f"request_id={req_id_cv.get()} instantiated OpenAI model_name={model_name}"
+            f"request_id={req_id_cv.get()} instantiated AI model_name={model_name}"
         )
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
@@ -47,7 +48,10 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
 
         try:
             LogWriter.info(f"request_id={req_id_cv.get()} ENTRY embed_documents()")
-            docs = self.embeddings.embed_documents(texts)
+            if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
+                docs = self.embeddings.embed_documents(texts, output_dimensionality=self.dimensions)
+            else:
+                docs = self.embeddings.embed_documents(texts)
             LogWriter.info(f"request_id={req_id_cv.get()} EXIT embed_documents()")
             metrics.llm_success_response_total.labels(self.model_name).inc()
             return docs
@@ -78,7 +82,10 @@ def embed_query(self, question: str) -> List[float]:
             logger.debug_pii(
                 f"request_id={req_id_cv.get()} embed_query() embedding question={question}"
             )
-            query_embedding = self.embeddings.embed_query(question)
+            if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
+                query_embedding = self.embeddings.embed_query(question, output_dimensionality=self.dimensions)
+            else:
+                query_embedding = self.embeddings.embed_query(question)
             LogWriter.info(f"request_id={req_id_cv.get()} EXIT embed_query()")
             metrics.llm_success_response_total.labels(self.model_name).inc()
             return query_embedding
@@ -105,8 +112,13 @@ async def aembed_query(self, question: str) -> List[float]:
         # metrics.llm_inprogress_requests.labels(self.model_name).inc()
 
         # try:
+        LogWriter.info(f"request_id={req_id_cv.get()} ENTRY aembed_query()")
         logger.debug_pii(f"aembed_query() embedding question={question}")
-        query_embedding = await self.embeddings.aembed_query(question)
+        if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
+            query_embedding = await self.embeddings.aembed_query(question, output_dimensionality=self.dimensions)
+        else:
+            query_embedding = await self.embeddings.aembed_query(question)
+        LogWriter.info(f"request_id={req_id_cv.get()} EXIT aembed_query()")
         # metrics.llm_success_response_total.labels(self.model_name).inc()
         return query_embedding
         # except Exception as e:
@@ -138,10 +150,8 @@ def __init__(self, config):
         super().__init__(
             config, model_name=config.get("model_name", "text-embedding-3-small")
         )
-        # from langchain_openai import OpenAIEmbeddings
-        from langchain_community.embeddings.openai import OpenAIEmbeddings
 
-        self.embeddings = OpenAIEmbeddings(model=self.model_name, base_url=self.base_url)
+        self.embeddings = OpenAIEmbeddings(model=self.model_name, base_url=config.get("base_url"))
 
 
 class VertexAI_PaLM_Embedding(EmbeddingModel):
@@ -151,7 +161,18 @@ def __init__(self, config):
         super().__init__(config, model_name=config.get("model_name", "VertexAI PaLM"))
         from langchain.embeddings import VertexAIEmbeddings
 
-        self.embeddings = VertexAIEmbeddings()
+        self.embeddings = VertexAIEmbeddings(model_name=self.model_name)
+
+
+class GenAI_Embedding(EmbeddingModel):
+    """Google GenAI Embedding Model"""
+
+    def __init__(self, config):
+        super().__init__(config, model_name=config.get("model_name", "gemini-embedding-exp-03-07"))
+
+        self.embeddings = GoogleGenerativeAIEmbeddings(model=self.model_name)
+        self.dimensions = config.get("dimensions", 1536)
+
 
 
 class AWS_Bedrock_Embedding(EmbeddingModel):

common/embeddings/tigergraph_embedding_store.py

@@ -277,7 +277,7 @@ def remove_embeddings(
         return
 
     def retrieve_similar(self, query_embedding, top_k=10, filter_expr: str = None, vertex_types: List[str] = ["DocumentChunk"]):
-        res = retrieve_similar_with_score(query_embedding, top_k=top_k, filter_expr=filter_expr, vertex_types=vertex_types)
+        res = self.retrieve_similar_with_score(query_embedding, top_k=top_k, filter_expr=filter_expr, vertex_types=vertex_types)
         similar = [x[0] for x in res]
         return similar
 

common/llm_services/__init__.py

@@ -3,8 +3,9 @@
 from .openai_service import OpenAI
 from .aws_sagemaker_endpoint import AWS_SageMaker_Endpoint
 from .google_vertexai_service import GoogleVertexAI
+from .google_genai_service import GoogleGenAI
 from .aws_bedrock_service import AWSBedrock
 from .groq_llm_service import Groq
 from .ollama import Ollama
 from .huggingface_endpoint import HuggingFaceEndpoint
-from .ibm_watsonx_service import IBMWatsonX
+from .ibm_watsonx_service import IBMWatsonX

common/llm_services/google_genai_service.py

@@ -0,0 +1,114 @@
+import logging
+import os
+
+from common.llm_services import LLM_Model
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+from common.logs.log import req_id_cv
+from common.logs.logwriter import LogWriter
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleGenAI(LLM_Model):
+    def __init__(self, config):
+        super().__init__(config)
+        for auth_detail in config["authentication_configuration"].keys():
+            os.environ[auth_detail] = config["authentication_configuration"][
+                auth_detail
+            ]
+
+        model_name = config["llm_model"]
+        self.llm = ChatGoogleGenerativeAI(
+            temperature=config["model_kwargs"]["temperature"],
+            model=model_name,
+            max_tokens=None,
+            timeout=None,
+            max_retries=2,
+        )
+        self.prompt_path = config["prompt_path"]
+        LogWriter.info(
+            f"request_id={req_id_cv.get()} instantiated OpenAI model_name={model_name}"
+        )
+
+    @property
+    def map_question_schema_prompt(self):
+        return self._read_prompt_file(self.prompt_path + "map_question_to_schema.txt")
+
+    @property
+    def generate_function_prompt(self):
+        return self._read_prompt_file(self.prompt_path + "generate_function.txt")
+
+    @property
+    def generate_cypher_prompt(self):
+        filepath = self.prompt_path + "generate_cypher.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().generate_cypher_prompt
+
+    @property
+    def entity_relationship_extraction_prompt(self):
+        return self._read_prompt_file(
+            self.prompt_path + "entity_relationship_extraction.txt"
+        )
+
+    @property
+    def route_response_prompt(self):
+        filepath = self.prompt_path + "route_response.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().route_response_prompt
+
+    @property
+    def graphrag_scoring_prompt(self):
+        filepath = self.prompt_path + "graphrag_scoring.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().graphrag_scoring_prompt
+
+    @property
+    def keyword_extraction_prompt(self):
+        filepath = self.prompt_path + "keyword_extraction.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().keyword_extraction_prompt
+
+    @property
+    def question_expansion_prompt(self):
+        filepath = self.prompt_path + "question_expansion.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().question_expansion_prompt
+
+    @property
+    def supportai_response_prompt(self):
+        filepath = self.prompt_path + "supportai_response.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().supportai_response_prompt
+
+    @property
+    def chatbot_response_prompt(self):
+        filepath = self.prompt_path + "chatbot_response.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().chatbot_response_prompt
+
+    @property
+    def hyde_prompt(self):
+        filepath = self.prompt_path + "hyde.txt"
+        if os.path.exists(filepath):
+            return self._read_prompt_file(filepath)
+        else:
+            return super().hyde_prompt
+
+    @property
+    def model(self):
+        return self.llm

common/llm_services/openai_service.py

@@ -1,10 +1,7 @@
 import logging
 import os
 
-if os.getenv("ECC"):
-    from langchain_openai.chat_models import ChatOpenAI
-else:
-    from langchain_community.chat_models import ChatOpenAI
+from langchain_openai.chat_models import ChatOpenAI
 
 from common.llm_services import LLM_Model
 from common.logs.log import req_id_cv

common/prompts/google_gemini/chatbot_response.txt

@@ -0,0 +1,32 @@
+You are a highly efficient and empathetic AI-powered customer support agent for TigerGraph. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines:
+
+1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary.
+2. **Provide Accurate Answers**: Use the given context in JSON format to deliver reliable and precise responses.
+3. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive.
+4. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions.
+5. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI.
+6. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support.
+7. **Reference Relevant Knowledge**: Ground your responses based on the context and resources. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID].
+8. **Here are the channels to reach the TigerGraph team**:
+   - [Contact Sales](https://www.tigergraph.com/contact/)
+   - [Support Portal](https://www.tigergraph.com/support/)
+   - [Documentation](https://docs.tigergraph.com/)
+   - [Community Forum](https://dev.tigergraph.com/)
+   - [Email](mailto:support@tigergraph.com)
+
+Example format for responses:
+- **Solution**: "Here's the information you requested: [details]."
+- **Clarification**: "Could you provide more details about [specific aspect]?"
+- **Escalation**: "This issue may require human assistance. Please reach out via [Support Portal](https://www.tigergraph.com/support/) or email us at [support@tigergraph.com](mailto:support@tigergraph.com)."
+
+Your mission is to ensure a seamless, satisfying customer experience while upholding TigerGraph's values and commitment to enterprise excellence.
+
+Give the context in JSON format, combine and rephrase it to answer the question.
+Use only the provided information in context without adding any reasoning or additional logic.
+Make sure all information in the context are covered in the generated answer.
+Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only.
+Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information.
+
+Question: {question}
+Context: {context}
+Format: {format_instructions}