Merge pull request #16 from igooch/messages

Adds the ability to ingest the optional field messages for the chat endpoint, which stores the history of the chat.

genai/api/genai_api/src/main.py

@@ -20,7 +20,8 @@
 import io
 import json
 import requests
-
+from typing import List
+from vertexai.language_models import ChatMessage
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -51,7 +52,7 @@
 
 
 app = FastAPI(
-    docs_url='/genai_docs', 
+    docs_url='/genai_docs',
     redoc_url=None,
     title="GenAI Quickstart APIs",
     description="Core APIs for the GenAI Quickstart for Gaming",
@@ -64,7 +65,7 @@
 )
 
 
-GENAI_GEMINI_ENDPOINT= os.environ['GENAI_GEMINI_ENDPOINT'] 
+GENAI_GEMINI_ENDPOINT= os.environ['GENAI_GEMINI_ENDPOINT']
 GENAI_TEXT_ENDPOINT  = os.environ['GENAI_TEXT_ENDPOINT']
 GENAI_CHAT_ENDPOINT  = os.environ['GENAI_CHAT_ENDPOINT']
 GENAI_CODE_ENDPOINT  = os.environ['GENAI_CODE_ENDPOINT']
@@ -102,6 +103,7 @@ class Payload_Vertex_Gemini(BaseModel):
 class Payload_Chat(BaseModel):
     prompt: str
     context: str | None = ''
+    message_history: List[ChatMessage] | None = []
     max_output_tokens: int | None = 1024
     temperature: float | None = 0.2
     top_p: float | None = 0.8
@@ -223,7 +225,7 @@ def genai_gemini(payload: Payload_Vertex_Gemini):
 def genai_text(payload: Payload_Text):
     try:
         request_payload = {
-            'prompt': payload.prompt, 
+            'prompt': payload.prompt,
             'max_output_tokens': payload.max_output_tokens,
             'temperature': payload.temperature,
             'top_p': payload.top_p,
@@ -246,6 +248,7 @@ def genai_chat(payload: Payload_Chat):
         request_payload = {
             'prompt': payload.prompt,
             'context': payload.context,
+            'message_history': payload.message_history,
             'max_output_tokens': payload.max_output_tokens,
             'temperature': payload.temperature,
             'top_p': payload.top_p,
@@ -266,7 +269,7 @@ def genai_chat(payload: Payload_Chat):
 def genai_code(payload: Payload_Code):
     try:
         request_payload = {
-            'prompt': payload.prompt, 
+            'prompt': payload.prompt,
             'max_output_tokens': payload.max_output_tokens,
             'temperature': payload.temperature,
             'top_p': payload.top_p,
@@ -287,7 +290,7 @@ def genai_code(payload: Payload_Code):
 def genai_image(payload: Payload_Image):
     try:
         request_payload = {
-            'prompt': payload.prompt, 
+            'prompt': payload.prompt,
             'number_of_images': payload.number_of_images,
             'seed': payload.seed,
         }

genai/api/genai_api/src/test_main.py

@@ -23,7 +23,7 @@
 
 @mock.patch('requests.post')
 def test_genai_text(mock_post):
-    
+
     # Define a mock response content as a JSON string
     expected_response = {'mocked_key': 'mocked_value'}
     mock_response_content = json.dumps(expected_response).encode()
@@ -56,7 +56,7 @@ def test_genai_text(mock_post):
 
 @mock.patch('requests.post')
 def test_genai_chat(mock_post):
-    
+
     # Define a mock response content as a JSON string
     expected_response = {'mocked_key': 'mocked_value'}
     mock_response_content = json.dumps(expected_response).encode()
@@ -73,6 +73,7 @@ def test_genai_chat(mock_post):
     payload = {
         "prompt": "test prompt",
         "context": "my context",
+        "message_history": [{"author": "me", "content": "my content"}],
         "max_output_tokens": 1024,
         "temperature": 0.2,
         "top_p": 0.8,
@@ -90,7 +91,7 @@ def test_genai_chat(mock_post):
 
 @mock.patch('requests.post')
 def test_genai_code(mock_post):
-    
+
     # Define a mock response content as a JSON string
     expected_response = {'mocked_key': 'mocked_value'}
     mock_response_content = json.dumps(expected_response).encode()
@@ -123,7 +124,7 @@ def test_genai_code(mock_post):
 
 @mock.patch('requests.post')
 def test_genai_image(mock_post):
-    
+
     # Define a mock response content as a JSON string
     expected_response = {'mocked_key': 'mocked_value'}
     mock_response_content = json.dumps(expected_response).encode()

genai/api/vertex_chat_api/src/main.py

@@ -14,13 +14,12 @@
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from fastapi.responses import StreamingResponse
 from utils.model_util import Google_Cloud_GenAI
-import io
-import os, sys
-import json
-import requests
+import sys
 import logging
+import requests
+from typing import List
+from vertexai.language_models import ChatMessage
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -68,6 +67,7 @@ def get_gcp_metadata():
 class Payload_Vertex_Chat(BaseModel):
     prompt: str
     context: str | None = ''
+    message_history: List[ChatMessage] | None = []
     max_output_tokens: int | None = 1024
     temperature: float | None = 0.2
     top_p: float | None = 0.8
@@ -88,6 +88,7 @@ def vertex_llm_chat(payload: Payload_Vertex_Chat):
         request_payload = {
             'prompt': payload.prompt,
             'context': payload.context,
+            'message_history': payload.message_history,
             'max_output_tokens': payload.max_output_tokens,
             'temperature': payload.temperature,
             'top_p': payload.top_p,

genai/api/vertex_chat_api/src/test_main.py

@@ -40,6 +40,7 @@ def test_genai(mock_post):
     payload = {
         "prompt": "test prompt",
         "context": "my test context",
+        "messages": [{"author": "me", "content": "my content"}],
         "max_output_tokens": 1024,
         "temperature": 0.2,
         "top_p": 0.8,

genai/api/vertex_chat_api/src/utils/model_util.py

@@ -25,20 +25,20 @@ class Google_Cloud_GenAI:
     def __init__(self, GCP_PROJECT_ID, GCP_REGION,  MODEL_TYPE):
         if GCP_PROJECT_ID=="":
             print(f'[ WARNING ] GCP_PROJECT_ID ENV variable is empty. Be sure to set the GCP_PROJECT_ID ENV variable.')
-        
+
         if GCP_REGION=="":
             print(f'[ WARNING ] GCP_REGION ENV variable is empty. Be sure to set the GCP_REGION ENV variable.')
 
         if MODEL_TYPE=="":
-            print(f'[ WARNING ] MODEL_TYPE ENV variable is empty. Be sure to set the MODEL_TYPE ENV variable.') 
-        
+            print(f'[ WARNING ] MODEL_TYPE ENV variable is empty. Be sure to set the MODEL_TYPE ENV variable.')
+
         self.GCP_PROJECT_ID = GCP_PROJECT_ID
         self.GCP_REGION = GCP_REGION
         self.MODEL_TYPE = MODEL_TYPE
         self.pretrained_model = f'{MODEL_TYPE.lower()}@001'
 
         self.vertexai = vertexai.init(project=GCP_PROJECT_ID, location=GCP_REGION)
-        
+
         if MODEL_TYPE.lower() == 'text-bison':
             self.model = TextGenerationModel.from_pretrained(self.pretrained_model)
         elif MODEL_TYPE.lower() == 'chat-bison':
@@ -52,7 +52,7 @@ def __init__(self, GCP_PROJECT_ID, GCP_REGION,  MODEL_TYPE):
             print(f'[ ERROR ] No MODEL_TYPE specified or MODEL_TYPE is incorrect. Expecting MODEL_TYPE ENV var of "text-bison", "chat-bison", "code-bison", or "codechat-bison".')
             sys.exit()
 
-    def call_llm(self, prompt, temperature=0.2, max_output_tokens=256, top_p=0.8, top_k=40, context='', chat_examples=[], code_suffix=''):
+    def call_llm(self, prompt, temperature=0.2, max_output_tokens=256, top_p=0.8, top_k=40, context='', chat_examples=[], message_history=[], code_suffix=''):
         if self.MODEL_TYPE.lower() == 'text-bison':
             try:
                 parameters = {
@@ -70,7 +70,7 @@ def call_llm(self, prompt, temperature=0.2, max_output_tokens=256, top_p=0.8, to
             except Exception as e:
                 print(f'[ EXCEPTION ] At call_llm for text-bison. {e}')
                 return ''
-        
+
         elif self.MODEL_TYPE.lower() == 'chat-bison':
             try:
                 '''
@@ -89,6 +89,7 @@ def call_llm(self, prompt, temperature=0.2, max_output_tokens=256, top_p=0.8, to
                 chat = self.model.start_chat(
                     context=context,
                     examples=chat_examples,
+                    message_history=message_history,
                     temperature=temperature,
                     max_output_tokens=max_output_tokens,
                     top_p=top_p,
@@ -100,16 +101,16 @@ def call_llm(self, prompt, temperature=0.2, max_output_tokens=256, top_p=0.8, to
             except Exception as e:
                 print(f'[ EXCEPTION ] At call_chat for chat-bison. {e}')
                 return ''
-        
+
         elif self.MODEL_TYPE.lower() == 'code-bison':
             '''A language model that generates code.'''
             try:
                 response = self.model.predict(prefix=prompt, temperature=temperature, max_output_tokens=max_output_tokens, suffix=code_suffix)
                 return response
             except Exception as e:
                 print(f'[ EXCEPTION ] At call_chat for codechat-bison. {e}')
-                return ''        
-        
+                return ''
+
         elif self.MODEL_TYPE.lower() == 'codechat-bison':
             '''CodeChatModel represents a model that is capable of completing code.'''
             try: