added logging

app.py

@@ -1,32 +1,64 @@
 import gradio as gr
 import weaviate
-from weaviate.embedded import EmbeddedOptions
 import os
 from openai import AsyncOpenAI
 from dotenv import load_dotenv
-import textwrap
 import asyncio
-import aiohttp
 from functools import wraps
+import logging
+import time
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 # Load environment variables
 load_dotenv()
 
 # Set up AsyncOpenAI client
 openai_client = AsyncOpenAI(api_key=os.getenv('OPENAI_API_KEY'))
 
-# Connect to Weaviate
-client = weaviate.Client(
-    url=os.getenv('WCS_URL'),
-    auth_client_secret=weaviate.auth.AuthApiKey(os.getenv('WCS_API_KEY')),
-    additional_headers={
-        "X-OpenAI-Api-Key": os.getenv('OPENAI_API_KEY')
-    }
-)
+# Initialize client as None
+client = None
 
 # Get the collection name from environment variable
 COLLECTION_NAME = os.getenv('WEAVIATE_COLLECTION_NAME')
 
+# Global variable to track connection status
+# Global variable to track connection status
+connection_status = {"status": "Disconnected", "color": "red"}
+
+# Function to initialize the Weaviate client
+async def initialize_weaviate_client(max_retries=3, retry_delay=5):
+    global client, connection_status
+    retries = 0
+    while retries < max_retries:
+        connection_status = {"status": "Connecting...", "color": "orange"}
+        try:
+            logger.info(f"Attempting to connect to Weaviate (Attempt {retries + 1}/{max_retries})")
+            client = weaviate.Client(
+                url=os.getenv('WCS_URL'),
+                auth_client_secret=weaviate.auth.AuthApiKey(os.getenv('WCS_API_KEY')),
+                additional_headers={
+                    "X-OpenAI-Api-Key": os.getenv('OPENAI_API_KEY')
+                }
+            )
+            # Test the connection
+            await asyncio.to_thread(client.schema.get)
+            connection_status = {"status": "Connected", "color": "green"}
+            logger.info("Successfully connected to Weaviate")
+            return connection_status
+        except Exception as e:
+            logger.error(f"Error connecting to Weaviate: {str(e)}")
+            connection_status = {"status": f"Error: {str(e)}", "color": "red"}
+            retries += 1
+            if retries < max_retries:
+                logger.info(f"Retrying in {retry_delay} seconds...")
+                await asyncio.sleep(retry_delay)
+            else:
+                logger.error("Max retries reached. Could not connect to Weaviate.")
+    return connection_status
+
+
 # Async-compatible caching decorator
 def async_lru_cache(maxsize=128):
     cache = {}
@@ -67,43 +99,41 @@ async def search_multimodal(query: str, limit: int = 30, alpha: float = 0.6):
         print(f"An error occurred during the search: {str(e)}")
         return []
 
-async def generate_response(query: str, context: str) -> str:
+async def generate_response_stream(query: str, context: str):
     prompt = f"""
 You are an AI assistant with extensive expertise in the semiconductor industry. Your knowledge spans a wide range of companies, technologies, and products, including but not limited to: System-on-Chip (SoC) designs, Field-Programmable Gate Arrays (FPGAs), Microcontrollers, Integrated Circuits (ICs), semiconductor manufacturing processes, and emerging technologies like quantum computing and neuromorphic chips.
 Use the following context, your vast knowledge, and the user's question to generate an accurate, comprehensive, and insightful answer. While formulating your response, follow these steps internally:
-
 Analyze the question to identify the main topic and specific information requested.
 Evaluate the provided context and identify relevant information.
 Retrieve additional relevant knowledge from your semiconductor industry expertise.
 Reason and formulate a response by combining context and knowledge.
 Generate a detailed response that covers all aspects of the query.
 Review and refine your answer for coherence and accuracy.
-
 In your output, provide only the final, polished response. Do not include your step-by-step reasoning or mention the process you followed.
 IMPORTANT: Ensure your response is grounded in factual information. Do not hallucinate or invent information. If you're unsure about any aspect of the answer or if the necessary information is not available in the provided context or your knowledge base, clearly state this uncertainty. It's better to admit lack of information than to provide inaccurate details.
 Your response should be:
-
 Thorough and directly address all aspects of the user's question
 Based solely on factual information from the provided context and your reliable knowledge
 Include specific examples, data points, or case studies only when you're certain of their accuracy
 Explain technical concepts clearly, considering the user may have varying levels of expertise
 Clearly indicate any areas where information is limited or uncertain
-
 Context: {context}
 User Question: {query}
 Based on the above context and your extensive knowledge of the semiconductor industry, provide your detailed, accurate, and grounded response below. Remember, only include information you're confident is correct, and clearly state any uncertainties: 
     """
 
-    response = await openai_client.chat.completions.create(
+    async for chunk in await openai_client.chat.completions.create(
         model="gpt-4o",
         messages=[
             {"role": "system", "content": "You are an expert Semi Conductor industry analyst"},
             {"role": "user", "content": prompt}
         ],
-        temperature=0
-    )
-
-    return response.choices[0].message.content
+        temperature=0,
+        stream=True
+    ):
+        content = chunk.choices[0].delta.content
+        if content is not None:
+            yield content
 
 def process_search_result(item):
     if item['content_type'] == 'text':
@@ -114,14 +144,12 @@ def process_search_result(item):
         return f"Table Description from {item['source_document']} (Page {item['page_number']}): {item['description']}\n\n"
     return ""
 
-async def esg_analysis(user_query: str):
+async def esg_analysis_stream(user_query: str):
     search_results = await search_multimodal(user_query)
     
     context_parts = await asyncio.gather(*[asyncio.to_thread(process_search_result, item) for item in search_results])
     context = "".join(context_parts)
     
-    response = await generate_response(user_query, context)
-
     sources = []
     for item in search_results[:5]:  # Limit to top 5 sources
         source = {
@@ -135,22 +163,9 @@ async def esg_analysis(user_query: str):
             source["image_path"] = item.get("image_path", "N/A")
         sources.append(source)
 
-    return response, sources
-
-def wrap_text(text, width=120):
-    return textwrap.fill(text, width=width)
-
-async def gradio_interface(user_question):
-    ai_response, sources = await esg_analysis(user_question)
-    
-    # Format AI response
-    formatted_response = f"""
-## AI Response
+    return generate_response_stream(user_query, context), sources
 
-{ai_response}
-    """
-    
-    # Format sources
+def format_sources(sources):
     source_text = "## Top 5 Sources\n\n"
     for i, source in enumerate(sources, 1):
         source_text += f"### Source {i}\n"
@@ -162,20 +177,154 @@ async def gradio_interface(user_question):
         if 'image_path' in source:
             source_text += f"- **Image Path:** {source['image_path']}\n"
         source_text += "\n"
+    return source_text
+
+# Custom CSS for the status box
+custom_css = """
+#status-box {
+    position: absolute;
+    top: 10px;
+    right: 10px;
+    background-color: white;
+    padding: 5px 10px;
+    border-radius: 5px;
+    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
+    z-index: 1000;
+    display: flex;
+    align-items: center;
+}
+#status-light {
+    width: 10px;
+    height: 10px;
+    border-radius: 50%;
+    display: inline-block;
+    margin-right: 5px;
+}
+#status-text {
+    font-size: 14px;
+    font-weight: bold;
+}
+"""
+
+def get_connection_status():
+    status = connection_status["status"]
+    color = connection_status["color"]
+    return f'<div id="status-box"><div id="status-light" style="background-color: {color};"></div><span id="status-text">{status}</span></div>'
+
+async def check_connection():
+    global connection_status
+    try:
+        if client:
+            await asyncio.to_thread(client.schema.get)
+            return {"status": "Connected", "color": "green"}
+        else:
+            return {"status": "Disconnected", "color": "red"}
+    except Exception:
+        return {"status": "Disconnected", "color": "red"}
+
+async def update_status():
+    global connection_status
+    while True:
+        new_status = await check_connection()
+        if new_status != connection_status:
+            connection_status = new_status
+            yield new_status
+        await asyncio.sleep(5)  # Check every 5 seconds
+
+async def gradio_interface(user_question):
+    if connection_status["status"] != "Connected":
+        return "Error: Database not connected. Please wait for the connection to be established.", ""
+
+    response_generator, sources = await esg_analysis_stream(user_question)
+    formatted_sources = format_sources(sources)
+    
+    full_response = ""
+    async for response_chunk in response_generator:
+        full_response += response_chunk
+    
+    return full_response, formatted_sources
+
+with gr.Blocks(css=custom_css) as iface:
+    status_indicator = gr.HTML(get_connection_status())
+    
+    with gr.Row():
+        gr.Markdown("# Semiconductor Industry Analysis")
+    
+    gr.Markdown("Ask questions about the semiconductor industry and get AI-powered answers with sources.")
+    
+    user_question = gr.Textbox(lines=2, placeholder="Enter your question about the semiconductor industry...", interactive=False)
+    ai_response = gr.Markdown(label="AI Response")
+    sources_output = gr.Markdown(label="Sources")
+    
+    submit_btn = gr.Button("Submit", interactive=False)
+    
+    submit_btn.click(
+        fn=gradio_interface,
+        inputs=user_question,
+        outputs=[ai_response, sources_output],
+    )
+
+    # Update status
+    def update_status_indicator(status):
+        return get_connection_status()  # Return the HTML string directly
+
+    def update_input_state(status):
+        is_connected = status["status"] == "Connected"
+        return gr.update(interactive=is_connected), gr.update(interactive=is_connected)
+
+    status_updater = gr.State(connection_status)
+    
+    iface.load(
+        lambda: connection_status,
+        outputs=[status_updater],
+        every=1,
+    )
+
+    status_updater.change(
+        fn=update_status_indicator,
+        inputs=[status_updater],
+        outputs=[status_indicator],
+    )
+
+    status_updater.change(
+        fn=update_input_state,
+        inputs=[status_updater],
+        outputs=[user_question, submit_btn],
+    )
+
+    status_updater = gr.State(connection_status)
+    
+    iface.load(
+        lambda: connection_status,
+        outputs=[status_updater],
+        every=1,
+    )
+
+    status_updater.change(
+        fn=update_status_indicator,
+        inputs=[status_updater],
+        outputs=[status_indicator],
+    )
+
+    status_updater.change(
+        fn=update_input_state,
+        inputs=[status_updater],
+        outputs=[user_question, submit_btn],
+    )
+
+async def main():
+    # Check environment variables
+    required_env_vars = ['WCS_URL', 'WCS_API_KEY', 'OPENAI_API_KEY', 'WEAVIATE_COLLECTION_NAME']
+    for var in required_env_vars:
+        if not os.getenv(var):
+            logger.error(f"Environment variable {var} is not set!")
+            return
+
+    # Initialize the client before launching the interface
+    await initialize_weaviate_client()
     
-    return formatted_response, source_text
-
-iface = gr.Interface(
-    fn=lambda user_question: asyncio.run(gradio_interface(user_question)),
-    inputs=gr.Textbox(lines=2, placeholder="Enter your question about the semiconductor industry..."),
-    outputs=[
-        gr.Markdown(label="AI Response"),
-        gr.Markdown(label="Sources")
-    ],
-    title="Semiconductor Industry ESG Analysis",
-    description="Ask questions about the semiconductor industry and get AI-powered answers with sources.",
-    flagging_dir="/app/flagged"  # Specify the flagging directory
-)
+    # Launch the interface regardless of connection status
+    await iface.launch(server_name="0.0.0.0", server_port=7860, share=True)
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    asyncio.run(main())