Update app.py

app.py

@@ -1,7 +1,3 @@
-# from langgraph.graph import Graph
-# from langchain_groq import ChatGroq
-# llm = langchain_groq(model="llama3-70b-8192")
-# llm.invoke("hi how are you")
 import streamlit as st
 import os
 import base64
@@ -15,60 +11,34 @@ from langchain.agents import Tool, initialize_agent
 from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
 from groq import Groq
 import open_clip
-from open_clip import create_model_from_pretrained, get_tokenizer # works on open-clip-torch>=2.23.0, timm>=0.9.8
+import torch
+from PIL import Image
 
-model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
-tokenizer = open_clip.get_tokenizer('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
+# Load environment variables
 load_dotenv()
 groq_api_key = os.getenv("GROQ_API_KEY")
-
 if not groq_api_key:
     st.error("Groq API Key not found in .env file")
     st.stop()
 
+# Configure Streamlit
 st.set_page_config(page_title="Medical Bot", page_icon="👨‍🔬")
 st.title("Medical Bot")
+
+# Initialize LLM models
 llm_text = ChatGroq(model="llama-3.3-70b-versatile", groq_api_key=groq_api_key)
 llm_image = ChatGroq(model="llama-3.2-90b-vision-preview", groq_api_key=groq_api_key)
 
-wikipedia_wrapper = WikipediaAPIWrapper()
-wikipedia_tool = Tool(
-    name="Wikipedia",
-    func=wikipedia_wrapper.run,
-    description="A tool for searching the Internet to find various information on the topics mentioned."
-)
-math_chain = LLMMathChain.from_llm(llm=llm_text)
-calculator = Tool(
-    name="Calculator",
-    func=math_chain.run,
-    description="A tool for solving mathematical problems. Provide only the mathematical expressions."
-)
-
-prompt = """
-You are a mathematical problem-solving assistant tasked with helping users solve their questions. Arrive at the solution logically, providing a clear and step-by-step explanation. Present your response in a structured point-wise format for better understanding.
-Question: {question}
-Answer:
-"""
-
-prompt_template = PromptTemplate(
-    input_variables=["question"],
-    template=prompt
-)
-# Combine all the tools into a chain for text questions
-chain = LLMChain(llm=llm_text, prompt=prompt_template)
+# Load BiomedCLIP model for image classification
+model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
+tokenizer = open_clip.get_tokenizer('hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224')
 
-reasoning_tool = Tool(
-    name="Reasoning Tool",
-    func=chain.run,
-    description="A tool for answering logic-based and reasoning questions."
-)
 def classify_image(image_path: str) -> str:
     """Classifies a medical image using BiomedCLIP."""
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     model.to(device).eval()
     
-    # Open and preprocess image
-    image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
+    image = preprocess_val(Image.open(image_path)).unsqueeze(0).to(device)
     labels = ["MRI scan", "X-ray", "histopathology", "CT scan", "ultrasound", "medical chart"]
     texts = tokenizer([f"this is a photo of {l}" for l in labels], context_length=256).to(device)
     
@@ -80,26 +50,43 @@ def classify_image(image_path: str) -> str:
     top_class = labels[sorted_indices[0][0].item()]
     return f"The image is classified as {top_class}."
 
-# Wrap BiomedCLIP as a LangChain tool
-biomed_clip_tool = Tool(
-    name="BiomedCLIP Image Classifier",
-    func=classify_image,
-    description="Classifies medical images into categories like MRI, X-ray, histopathology, etc."
-)
-# Initialize the agents for text questions
+# Define tools
+wikipedia_tool = Tool(name="Wikipedia", func=WikipediaAPIWrapper().run, description="A tool for searching information.")
+math_chain = LLMMathChain.from_llm(llm=llm_text)
+calculator = Tool(name="Calculator", func=math_chain.run, description="Solves mathematical problems.")
+
+prompt_template = PromptTemplate(input_variables=["question"], template="""
+You are a mathematical problem-solving assistant. Solve the question step by step.
+Question: {question}
+Answer:
+""")
+chain = LLMChain(llm=llm_text, prompt=prompt_template)
+reasoning_tool = Tool(name="Reasoning Tool", func=chain.run, description="Answers logic-based questions.")
+
+biomed_clip_tool = Tool(name="BiomedCLIP Image Classifier", func=classify_image, description="Classifies medical images.")
+
+# Initialize agents
 assistant_agent_text = initialize_agent(
-    tools=[wikipedia_tool, calculator, reasoning_tool, biomed_clip_tool],
+    tools=[wikipedia_tool, calculator, reasoning_tool],
     llm=llm_text,
     agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
     verbose=False,
     handle_parsing_errors=True
 )
 
+assistant_agent_image = initialize_agent(
+    tools=[biomed_clip_tool],
+    llm=llm_image,
+    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+    verbose=False,
+    handle_parsing_errors=True
+)
+
+# Streamlit session state for chat messages
 if "messages" not in st.session_state:
-    st.session_state["messages"] = [
-        {"role": "assistant", "content": "Welcome! I am your Assistant. How can I help you today?"}
-    ]
+    st.session_state["messages"] = [{"role": "assistant", "content": "Welcome! How can I help you today?"}]
 
+# Chat Interface
 for msg in st.session_state.messages:
     if msg["role"] == "user" and "image" in msg:
         st.chat_message(msg["role"]).write(msg['content'])
@@ -112,87 +99,48 @@ if st.sidebar.button("Text Question"):
     st.session_state["section"] = "text"
 if st.sidebar.button("Image Question"):
     st.session_state["section"] = "image"
-
 if "section" not in st.session_state:
     st.session_state["section"] = "text"
 
 def clean_response(response):
-    if "```" in response:
-        response = response.split("```")[1].strip()
-    return response
+    return response.split("```")[-1].strip() if "```" in response else response
 
 if st.session_state["section"] == "text":
     st.header("Text Question")
-    st.write("Please enter your question below, and I will provide a detailed description of the problem and suggest a solution for it.")
     question = st.text_area("Your Question:")
     if st.button("Get Answer"):
         if question:
             with st.spinner("Generating response..."):
                 st.session_state.messages.append({"role": "user", "content": question})
                 st.chat_message("user").write(question)
-
-                st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=False)
-                try:
-                    response = assistant_agent_text.run(st.session_state.messages, callbacks=[st_cb])
-                    cleaned_response = clean_response(response)
-                    st.session_state.messages.append({'role': 'assistant', "content": cleaned_response})
-                    st.write('### Response:')
-                    st.success(cleaned_response)
-                except ValueError as e:
-                    st.error(f"An error occurred: {e}")
+                
+                response = assistant_agent_text.run(question)
+                cleaned_response = clean_response(response)
+                st.session_state.messages.append({'role': 'assistant', "content": cleaned_response})
+                st.write('### Response:')
+                st.success(cleaned_response)
         else:
-            st.warning("Please enter a question to get an answer.")
+            st.warning("Please enter a question.")
 
 elif st.session_state["section"] == "image":
     st.header("Image Question")
-    st.write("Please enter your question below and upload the medical image. I will provide a detailed description of the problem and suggest a solution for it.")
-    question = st.text_area("Your Question:", "Example: What is the patient suffering from?")
+    question = st.text_area("Your Question:")
     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-
     if st.button("Get Answer"):
-        if question and uploaded_file is not None:
+        if question and uploaded_file:
             with st.spinner("Generating response..."):
-                image_data = uploaded_file.read()
-                image_data_url = f"data:image/jpeg;base64,{base64.b64encode(image_data).decode()}"
-                st.session_state.messages.append({"role": "user", "content": question, "image": image_data})
+                image_path = f"temp_{uploaded_file.name}"
+                with open(image_path, "wb") as f:
+                    f.write(uploaded_file.read())
+                
+                st.session_state.messages.append({"role": "user", "content": question, "image": image_path})
                 st.chat_message("user").write(question)
-                st.image(image_data, caption='Uploaded Image', use_column_width=True)
-
-                client = Groq()
-
-                messages = [
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": question
-                            },
-                            {
-                                "type": "image_url",
-                                "image_url": {
-                                    "url": image_data_url
-                                }
-                            }
-                        ]
-                    }
-                ]
-                try:
-                    completion = client.chat.completions.create(
-                        model="llama-3.2-90b-vision-preview",
-                        messages=messages,
-                        temperature=1,
-                        max_tokens=1024,
-                        top_p=1,
-                        stream=False,
-                        stop=None,
-                    )
-                    response = completion.choices[0].message.content
-                    cleaned_response = clean_response(response)
-                    st.session_state.messages.append({'role': 'assistant', "content": cleaned_response})
-                    st.write('### Response:')
-                    st.success(cleaned_response)
-                except ValueError as e:
-                    st.error(f"An error occurred: {e}")
+                st.image(image_path, caption='Uploaded Image', use_column_width=True)
+                
+                response = assistant_agent_image.run(image_path)
+                cleaned_response = clean_response(response)
+                st.session_state.messages.append({'role': 'assistant', "content": cleaned_response})
+                st.write('### Response:')
+                st.success(cleaned_response)
         else:
-            st.warning("Please enter a question and upload an image to get an answer.")
+            st.warning("Please enter a question and upload an image.")