added databricks model

app_config.py

@@ -1,12 +1,16 @@
 ISSUES = ['Anxiety','Suicide']
-SOURCES = ['OA_rolemodel', 'OA_finetuned']
+SOURCES = ['OA_rolemodel', 
+        #    'OA_finetuned',
+           "CTL_llama2"]
 SOURCES_LAB = {"OA_rolemodel":'OpenAI GPT3.5',
-               "OA_finetuned":'Finetuned OpenAI'}
+               "OA_finetuned":'Finetuned OpenAI',
+               "CTL_llama2": "Custom CTL"
+               }
 
 def source2label(source):
     return SOURCES_LAB[source]
 
-ENVIRON = "prod"
+ENVIRON = "dev"
 
 DB_SCHEMA = 'prod_db' if ENVIRON == 'prod' else 'test_db' 
 DB_CONVOS = 'conversations'

app_utils.py

@@ -7,6 +7,7 @@ from langchain.memory import ConversationBufferMemory
 from app_config import ENVIRON
 from models.openai.finetuned_models import finetuned_models, get_finetuned_chain
 from models.openai.role_models import get_role_chain, role_templates
+from models.databricks.scenario_sim_biz import get_databricks_chain
 from mongo_utils import new_convo
 
 langchain.verbose = ENVIRON=="dev"
@@ -35,7 +36,7 @@ def change_memories(memories, username, language, changed_source=False):
         if (memory not in st.session_state) or changed_source:
             source = params['source']
             logger.info(f"Source for memory {memory} is {source}")
-            if source in ('OA_rolemodel','OA_finetuned'):
+            if source in ('OA_rolemodel','OA_finetuned',"CTL_llama2"):
                 st.session_state[memory] = ConversationBufferMemory(ai_prefix='texter', human_prefix='helper')
     
     if ("convo_id" in st.session_state) and changed_source:
@@ -63,4 +64,11 @@ def get_chain(issue, language, source, memory, temperature):
         return get_finetuned_chain(OA_engine, memory, temperature)
     elif source in ('OA_rolemodel'):
         template = role_templates[f"{issue}-{language}"]
-        return get_role_chain(template, memory, temperature)
+        return get_role_chain(template, memory, temperature)
+    elif source in ('CTL_llama2'):
+        if language == "English":
+            language = "en"
+        elif language == "Spanish":
+            language = "es"
+        return get_databricks_chain(issue, language, memory, temperature)
+    

models/databricks/scenario_sim_biz.py

@@ -0,0 +1,66 @@
+import os
+import json
+import requests
+import logging
+from models.custom_parsers import CustomStringOutputParser
+from langchain.chains import ConversationChain
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.language_models.llms import LLM
+from langchain.prompts import PromptTemplate
+
+from typing import Any, List, Mapping, Optional, Dict
+
+class DatabricksCustomLLM(LLM):
+    issue:str
+    language:str
+    temperature:float = 0.8
+    db_url:str = os.environ['DATABRICKS_URL'] 
+    headers:Mapping[str,str] = {'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}', 'Content-Type': 'application/json'}
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom_databricks"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        data_ = {'inputs': {
+            'prompt': [prompt],
+            'issue': [self.issue],
+            'language': [self.language],
+            'temperature': [self.temperature]
+        }}
+        data_json = json.dumps(data_, allow_nan=True)
+        response = requests.request(method='POST', headers=self.headers, url=self.db_url, data=data_json)
+        
+        if response.status_code != 200:
+            raise Exception(f'Request failed with status {response.status_code}, {response.text}')
+        return response.json()["predictions"][0]["generated_text"]
+
+_DATABRICKS_TEMPLATE_ = """{history}
+helper: {input}
+texter:"""
+
+def get_databricks_chain(issue, language, memory, temperature=0.8):
+
+    PROMPT = PromptTemplate(
+        input_variables=['history', 'input'],
+        template=_DATABRICKS_TEMPLATE_
+    )
+    llm = DatabricksCustomLLM(
+        issue=issue,
+        language=language,
+        temperature=temperature
+    )
+    llm_chain = ConversationChain(
+        llm=llm,
+        prompt=PROMPT,
+        memory=memory,
+        output_parser=CustomStringOutputParser()
+    )
+    logging.debug(f"loaded Databricks Scenario Sim model")
+    return llm_chain, "helper:"