v0.1.1

src/llm/geminiLLM.py

@@ -0,0 +1,114 @@
+from typing import Any, List, Mapping, Optional, Dict 
+from pydantic import Extra, Field #, root_validator, model_validator
+import os,json
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.base import LLM
+import google.generativeai as genai
+from google.generativeai import types
+import ast
+#from langchain.llms import GooglePalm
+import requests,logging
+
+logger=logging.getLogger("llm")
+
+class GeminiLLM(LLM):
+    
+    model_name: str = "gemini-pro"
+    temperature: float = 0
+    max_tokens: int = 2048
+    stop: Optional[List] = []
+    prev_prompt: Optional[str]=""
+    prev_stop: Optional[str]=""
+    prev_run_manager:Optional[Any]=None
+    model: Optional[Any]=None
+
+    def __init__(
+        self,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.model=genai.GenerativeModel(self.model_name)
+        #self.model = palm.Text2Text(self.model_name)
+
+    @property
+    def _llm_type(self) -> str:
+        return "text2text-generation"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+    ) -> str:
+        self.prev_prompt=prompt
+        self.prev_stop=stop
+        self.prev_run_manager=run_manager
+        #print(types.SafetySettingDict)
+        if stop == None:
+            stop=self.stop
+        logger.debug("\nLLM in use is:" +self._llm_type)
+        logger.debug("Request to LLM is "+prompt)
+        
+        response=self.model.generate_content(prompt,
+                                    generation_config={"stop_sequences":self.stop,
+                                    "temperature":self.temperature, "max_output_tokens":self.max_tokens},
+                            safety_settings=[{"category":"HARM_CATEGORY_SEXUALLY_EXPLICIT","threshold":"BLOCK_NONE"},
+                             {"category":"HARM_CATEGORY_HATE_SPEECH","threshold":"BLOCK_NONE"},
+                             {"category":"HARM_CATEGORY_HARASSMENT","threshold":"BLOCK_NONE"},
+                             {"category":"HARM_CATEGORY_DANGEROUS_CONTENT","threshold":"BLOCK_NONE"}],
+                             stream=False
+                        )
+        try:
+            val=response.text
+            if val == None:
+                logger.debug("Response from LLM was None\n")
+                filterStr=""
+                for item in response.filters:
+                    for key,val in item.items():
+                        filterStr+=key+":"+str(val)
+                logger.error("Will switch to fallback LLM as response from palm is None::"+filterStr)
+                raise(Exception)
+            else:
+                logger.debug("Response from LLM "+val)
+        except Exception as ex:
+                logger.error("Will switch to fallback LLM as response from palm is None::")
+                raise(Exception)
+        if run_manager:
+            run_manager.on_llm_end(val)
+        return val
+    
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {"name": self.model_name, "type": "palm"} 
+    
+    def extractJson(self,val:str) -> Any:
+        """Helper function to extract json from this LLMs output"""
+        #This is assuming the json is the first item within ````
+        # palm is responding always with ```json and ending with ```, however sometimes response is not complete
+        # in case trailing ``` is not seen, we will call generation again with prev_prompt and result appended to it
+        try:
+            count=0
+            while val.startswith("```json") and not val.endswith("```") and count<7:
+                val=self._call(prompt=self.prev_prompt+" "+val,stop=self.prev_stop,run_manager=self.prev_run_manager)
+                count+=1
+            v2=val.replace("```json","```").split("```")[1]
+            try:
+                v4=json.loads(v2)
+            except:
+                #v3=v2.replace("\n","").replace("\r","").replace("'","\"")
+                v3=json.dumps(ast.literal_eval(v2))
+                v4=json.loads(v3)
+        except:
+            v2=val.replace("\n","").replace("\r","")
+            v3=json.dumps(ast.literal_eval(val))
+            #v3=v2.replace("'","\"")
+            v4=json.loads(v3)
+            #v4=json.loads(v2)
+        return v4
+    
+    def extractPython(self,val:str) -> Any:
+        """Helper function to extract python from this LLMs output"""
+        #This is assuming the python is the first item within ````
+        v2=val.replace("```python","```").split("```")[1]
+        return v2

src/llm/llmFactory.py

@@ -3,6 +3,7 @@ from baseInfra.dbInterface import DbInterface
 from llm.hostedLLM import HostedLLM
 from llm.togetherLLM import TogetherLLM
 from llm.palmLLM import PalmLLM
+from llm.geminiLLM import GeminiLLM
 
 
 class LLMFactory:
@@ -49,6 +50,8 @@ class LLMFactory:
             return TogetherLLM(**llm_config)
         elif llm_type == "palmLLM":
             return PalmLLM(**llm_config)
+        elif llm_type  == "geminiLLM":
+            return GeminiLLM(**llm_config)
         else:
             logger.error(f"Invalid LLM type: {llm_type}")
             raise ValueError(f"Invalid LLM type: {llm_type}")

src/main.py

@@ -3,7 +3,7 @@ import logging,os
 import fastapi
 from fastapi import Body, Depends
 import uvicorn
-from fastapi import HTTPException , status
+from fastapi import BackgroundTasks,HTTPException , status
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import FastAPI as Response
@@ -60,7 +60,7 @@ app.add_middleware(
 api_base="/api/v1"
 
 @app.post(api_base+"/getMatchingDocs")
-async def get_matching_docs(inStr: str, kwargs: Dict [Any, Any] ) -> Any:
+async def get_matching_docs(inStr: str, kwargs: Dict [Any, Any] ,background_tasks:BackgroundTasks) -> Any:
     """
     Gets the query embeddings and uses metadata appropriately and gets the matching docs for query
     TODO: Add parameter for type of query and number of docs to return