Added Llama3 support

app.py

@@ -27,23 +27,23 @@ def main():
     if analysis_type == "Code review":
       res=review_code(code, prompt)
       #st.text_area(label="Result",value=res, height=300)
-      st.markdown(res)
+      st.markdown(res['text'])
       time.sleep(1)
       #st.markdown("Hello ooo")
 
     elif analysis_type == "Code refinement":
       res=refine_code(code, prompt)
       #st.text_area(label="Result",value=res, height=300)
-      st.markdown(res)
+      st.markdown(res['text'])
 
     elif analysis_type == "Documentation":
       res=generate_documentation(code, prompt)
       #st.text_area(label="Result",value=res, height=300)
-      st.markdown(res)
+      st.markdown(res['text'])
     elif analysis_type == "Resume Writer":
       res=resume_writer(code, prompt)
       #st.text_area(label="Result",value=res, height=300)
-      st.markdown(res)
+      st.markdown(res['text'])
   
     st.success(f"Code analysis for {analysis_type} submitted successfully!")  
 

gemini.py

@@ -35,11 +35,12 @@ class GeminiModel:
     def execute(self, prompt: str) -> str:
         
         try:
-            total_tokens = self.model.count_tokens(prompt).total_tokens
+            prompt_tokens = self.model.count_tokens(prompt).total_tokens
             print(f"Input tokens: {total_tokens}")
             response = self.model.generate_content(prompt, generation_config=generation_config)
             output_tokens = self.model.count_tokens(response.text).total_tokens
             print(f"Output tokens: {output_tokens}")
-            return response.text
+
+            return response.text,{'prompt_tokens':prompt_tokens,"total_tokens":output_tokens}
         except Exception as e:
             return f"An error occurred: {e}"

handler.py

@@ -1,5 +1,8 @@
 import gemini
+import llama_groq
 from prompts import *
+#model = gemini.GeminiModel()
+model=llama_groq.LlamaModel()
 
 def review_code(code, c_prompt=None):
     if code is None or len(code) < 5 or code.isspace():
@@ -11,7 +14,7 @@ def review_code(code, c_prompt=None):
         #prompt = validation_prompt(code.strip())
         prompt = default_review_prompt1(code.strip())
 
-    model = gemini.GeminiModel()
+    
     try:
         res = model.execute(prompt)
     except Exception as e:
@@ -29,7 +32,7 @@ def refine_code(code, c_prompt=None):
         #prompt = validation_prompt(code.strip())
         prompt = default_refine_prompt(code.strip())
 
-    model = gemini.GeminiModel()
+    
     try:
         res = model.execute(prompt)
     except Exception as e:
@@ -47,7 +50,7 @@ def generate_documentation(code,c_prompt):
     else:
         prompt = default_doc_prompt(code.strip())
 
-    model = gemini.GeminiModel()
+    
     try:
         res = model.execute(prompt)
     except Exception as e:
@@ -65,7 +68,7 @@ def resume_writer(code,c_prompt):
     else:
         prompt = resume_prompt(code.strip())
 
-    model = gemini.GeminiModel()
+    
     try:
         res = model.execute(prompt)
     except Exception as e:

llama_groq.py

@@ -30,11 +30,9 @@ class LlamaModel:
     def execute(self, prompt: str) -> str:
         
         try:
-            #total_tokens = self.model.count_tokens(prompt).total_tokens
-            #print(f"Input tokens: {total_tokens}")
             response = self.model.invoke(prompt)
-            #output_tokens = response
-            #print(f"Output tokens: {output_tokens}")
-            return response
+            res=response.content
+            meta=response.response_metadata
+            return {'text':res,'meta':meta}
         except Exception as e:
             return f"An error occurred: {e}"