shifted to google gemini

TechdocsAPI/backend/__init__.py

@@ -1,5 +1,4 @@
 import mysql.connector
-from mysql.connector import errorcode
 
 from fastapi import FastAPI, status
 from fastapi.exceptions import HTTPException
@@ -8,14 +7,16 @@ from fastapi.templating import Jinja2Templates
 from backend.utils import DBConnection
 from backend.core.ConfigEnv import config
 
-from langchain.llms import Clarifai
+# from langchain.llms import Clarifai
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
+from langchain_google_genai import GoogleGenerativeAI
 
-app = FastAPI(title="Techdocs",
-              version="V0.0.1",
-              description="API for automatic code documentation generation!"
-              )
+app = FastAPI(
+    title="Techdocs",
+    version="V0.0.1",
+    description="API for automatic code documentation generation!",
+)
 
 from backend import router
 
@@ -24,31 +25,23 @@ try:
     test_conn = DBConnection.get_client().get_server_info()
 
     # send prompt wizardcoderLM-70b-instruct-GGUF model
-    with open("backend/utils/prompt.txt",'r') as f:
+    with open("backend/utils/prompt.txt", "r") as f:
         prompt = f.read()
 
-    prompt = PromptTemplate(template=prompt, input_variables=['instruction'])
+    prompt = PromptTemplate(template=prompt, input_variables=["instruction"])
 
-    llm = Clarifai(
-        pat = config.CLARIFAI_PAT,
-        user_id = config.USER_ID,
-        app_id = config.APP_ID, 
-        model_id = config.MODEL_ID,
-        model_version_id=config.MODEL_VERSION_ID,
+    llm = GoogleGenerativeAI(
+        model = "gemini-pro",
+        google_api_key=config.GOOGLE_API_KEY,
     )
 
-    llmchain = LLMChain(
-        prompt=prompt,
-        llm=llm
-    )
+    llmchain = LLMChain(prompt=prompt, llm=llm)
     app.state.llmchain = llmchain
 
-
     app.state.templates = Jinja2Templates(directory="./backend/templates")
 
 
-
 except mysql.connector.Error as err:
-    raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(err))    
-
-
+    raise HTTPException(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(err)
+    )

TechdocsAPI/backend/core/ConfigEnv.py

@@ -14,15 +14,9 @@ class Settings(BaseSettings):
     JWT_SECRET_KEY:str
     JWT_REFRESH_SECRET_KEY:str
     JWT_VERIFICATION_SECRET_KEY:str
-    # OPENAI_KEY:str
-    APP_ID:str
-    USER_ID:str
-    MODEL_ID:str
-    CLARIFAI_PAT:str
-    MODEL_VERSION_ID:str
+    GOOGLE_API_KEY:str
 
     MAIL_SERVER_URL:str
-    
     MAIL_USERNAME:str
     MAIL_PASSWORD:str
     MAIL_FROM:str

TechdocsAPI/backend/services/auth/ops.py

@@ -128,9 +128,9 @@ def ops_inference(source_code:str,api_key:str,username:str):
     def generate_docstring(source_code_message: str):
 
 
-        llm_response = app.state.llmchain.run({"instruction": source_code_message})
+        llm_response = app.state.llmchain.invoke({"instruction": source_code_message})
 
-        docstring = Inference(docstr=llm_response)        
+        docstring = Inference(docstr=llm_response["text"])        
     
         return docstring
 

TechdocsAPI/requirements.txt

@@ -5,8 +5,8 @@ pydantic==1.10.12
 python-jose[cryptography]
 passlib[bcrypt]
 mysql-connector-python
+langchain-google-genai
 pydantic[email]
 langchain
-clarifai
 Pillow
 jinja2