Update app.py

app.py

@@ -1,178 +1,191 @@
-# app.py
-import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
-from diffusers import StableDiffusionPipeline, DiffusionPipeline
-import torch
-from PIL import Image
-import numpy as np
+import anthropic
+import base64
+import json
 import os
-import tempfile
-import moviepy.editor as mpe
-import nltk
-from pydub import AudioSegment
-import warnings
-import asyncio
-import edge_tts
-import random
-from openai import OpenAI
-
-warnings.filterwarnings("ignore", category=UserWarning)
-
-# Ensure NLTK data is downloaded
-nltk.download('punkt')
-
-# LLM Inference Class
-class LLMInferenceNode:
-    def __init__(self):
-        self.huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-        self.sambanova_api_key = os.getenv("SAMBANOVA_API_KEY")
-        
-        self.huggingface_client = OpenAI(
-            base_url="https://api-inference.huggingface.co/v1/",
-            api_key=self.huggingface_token,
-        )
-        self.sambanova_client = OpenAI(
-            api_key=self.sambanova_api_key,
-            base_url="https://api.sambanova.ai/v1",
-        )
-
-    def generate(self, input_text, long_talk=True, compress=False, 
-                compression_level="medium", poster=False, prompt_type="Short",
-                provider="Hugging Face", model=None):
-        try:
-            # Define system message
-            system_message = "You are a helpful assistant. Try your best to give the best response possible to the user."
-            
-            # Define base prompts based on type
-            prompts = {
-                "Short": """Create a brief, straightforward caption for this description, suitable for a text-to-image AI system. 
-                        Focus on the main elements, key characters, and overall scene without elaborate details.""",
-                "Long": """Create a detailed visually descriptive caption of this description for a text-to-image AI system. 
-                        Include detailed visual descriptions, cinematography, and lighting setup."""
-            }
-            
-            base_prompt = prompts.get(prompt_type, prompts["Short"])
-            user_message = f"{base_prompt}\nDescription: {input_text}"
-
-            # Generate with selected provider
-            if provider == "Hugging Face":
-                client = self.huggingface_client
-            else:
-                client = self.sambanova_client
-
-            response = client.chat.completions.create(
-                model=model or "meta-llama/Meta-Llama-3.1-70B-Instruct",
-                max_tokens=1024,
-                temperature=1.0,
-                messages=[
-                    {"role": "system", "content": system_message},
-                    {"role": "user", "content": user_message},
-                ]
-            )
-            
-            return response.choices[0].message.content.strip()
-
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            return f"Error occurred while processing the request: {str(e)}"
-
-# Initialize models
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
-
-# Story generator
-story_generator = pipeline(
-    'text-generation',
-    model='gpt2-large',
-    device=0 if device == 'cuda' else -1
-)
-
-# Stable Diffusion model
-sd_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch_dtype
-).to(device)
-
-# Text-to-Speech function using edge_tts
-async def _text2speech_async(text):
-    communicate = edge_tts.Communicate(text, voice="en-US-AriaNeural")
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
-        tmp_path = tmp_file.name
-        await communicate.save(tmp_path)
-    return tmp_path
-
-def text2speech(text):
+import pandas as pd
+import pytz
+import re
+import streamlit as st
+from datetime import datetime
+from gradio_client import Client
+from azure.cosmos import CosmosClient, exceptions
+
+# App Configuration
+title = "🤖 ArXiv and Claude AI Assistant"
+st.set_page_config(page_title=title, layout="wide")
+
+# Cosmos DB configuration
+ENDPOINT = "https://acae-afd.documents.azure.com:443/"
+Key = os.environ.get("Key")
+DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
+CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME")
+
+# Initialize Anthropic client
+anthropic_client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
+
+# Initialize session state
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+def generate_filename(prompt, file_type):
+    """Generate a filename with timestamp and sanitized prompt"""
+    central = pytz.timezone('US/Central')
+    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
+    safe_prompt = re.sub(r'\W+', '', prompt)[:90]
+    return f"{safe_date_time}{safe_prompt}.{file_type}"
+
+def create_file(filename, prompt, response, should_save=True):
+    """Create and save a file with prompt and response"""
+    if not should_save:
+        return
+    with open(filename, 'w', encoding='utf-8') as file:
+        file.write(f"Prompt:\n{prompt}\n\nResponse:\n{response}")
+
+def save_to_cosmos_db(container, query, response1, response2):
+    """Save interaction to Cosmos DB"""
     try:
-        output_path = asyncio.run(_text2speech_async(text))
-        return output_path
+        if container:
+            timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S%f')
+            record = {
+                "id": timestamp,
+                "name": timestamp,
+                "query": query,
+                "response1": response1,
+                "response2": response2,
+                "timestamp": datetime.utcnow().isoformat(),
+                "type": "ai_response",
+                "version": "1.0"
+            }
+            container.create_item(body=record)
+            st.success(f"Record saved to Cosmos DB with ID: {record['id']}")
     except Exception as e:
-        print(f"Error in text2speech: {str(e)}")
-        raise
-
-def generate_story(prompt):
-    generated = story_generator(prompt, max_length=500, num_return_sequences=1)
-    story = generated[0]['generated_text']
-    return story
-
-def split_story_into_sentences(story):
-    sentences = nltk.sent_tokenize(story)
-    return sentences
-
-def generate_images(sentences):
-    images = []
-    for idx, sentence in enumerate(sentences):
-        image = sd_pipe(sentence).images[0]
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f"_{idx}.png")
-        image.save(temp_file.name)
-        images.append(temp_file.name)
-    return images
-
-def generate_audio(story_text):
-    audio_path = text2speech(story_text)
-    audio = AudioSegment.from_file(audio_path)
-    total_duration = len(audio) / 1000
-    return audio_path, total_duration
+        st.error(f"Error saving to Cosmos DB: {str(e)}")
 
-def compute_sentence_durations(sentences, total_duration):
-    total_words = sum(len(sentence.split()) for sentence in sentences)
-    return [total_duration * (len(sentence.split()) / total_words) for sentence in sentences]
-
-def create_video(images, durations, audio_path):
-    clips = [mpe.ImageClip(img).set_duration(dur) for img, dur in zip(images, durations)]
-    video = mpe.concatenate_videoclips(clips, method='compose')
-    audio = mpe.AudioFileClip(audio_path)
-    video = video.set_audio(audio)
-    output_path = os.path.join(tempfile.gettempdir(), "final_video.mp4")
-    video.write_videofile(output_path, fps=1, codec='libx264')
-    return output_path
-
-def process_pipeline(prompt, progress=gr.Progress()):
+def search_arxiv(query):
+    """Search ArXiv using Gradio client"""
     try:
-        story = generate_story(prompt)
-        sentences = split_story_into_sentences(story)
-        images = generate_images(sentences)
-        audio_path, total_duration = generate_audio(story)
-        durations = compute_sentence_durations(sentences, total_duration)
-        video_path = create_video(images, durations, audio_path)
-        return video_path
+        client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
+        
+        # Get response from Mixtral model
+        result_mixtral = client.predict(
+            query,
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            True,
+            api_name="/ask_llm"
+        )
+        
+        # Get response from Mistral model
+        result_mistral = client.predict(
+            query,
+            "mistralai/Mistral-7B-Instruct-v0.2",
+            True,
+            api_name="/ask_llm"
+        )
+        
+        # Get RAG-enhanced response
+        result_rag = client.predict(
+            query,
+            10,  # llm_results_use
+            "Semantic Search",
+            "mistralai/Mistral-7B-Instruct-v0.2",
+            api_name="/update_with_rag_md"
+        )
+        
+        return result_mixtral, result_mistral, result_rag
     except Exception as e:
-        print(f"Error in process_pipeline: {str(e)}")
-        raise gr.Error(f"An error occurred: {str(e)}")
-
-# Gradio Interface
-title = """<h1 align="center">AI Story Video Generator 🎥</h1>
-<p align="center">Generate a story from a prompt, create images for each sentence, and produce a video with narration!</p>"""
+        st.error(f"Error searching ArXiv: {str(e)}")
+        return None, None, None
 
-with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
-    gr.HTML(title)
-    
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(label="Enter a Prompt", lines=2)
-            generate_button = gr.Button("Generate Video")
-        with gr.Column():
-            video_output = gr.Video(label="Generated Video")
-    
-    generate_button.click(fn=process_pipeline, inputs=prompt_input, outputs=video_output)
+def main():
+    st.title(title)
 
-demo.launch(debug=True)
+    # Initialize Cosmos DB client if key is available
+    if Key:
+        cosmos_client = CosmosClient(ENDPOINT, credential=Key)
+        try:
+            database = cosmos_client.get_database_client(DATABASE_NAME)
+            container = database.get_container_client(CONTAINER_NAME)
+        except Exception as e:
+            st.error(f"Error connecting to Cosmos DB: {str(e)}")
+            container = None
+    else:
+        st.warning("Cosmos DB Key not found in environment variables")
+        container = None
+
+    # Create tabs for different functionalities
+    arxiv_tab, claude_tab, history_tab = st.tabs(["ArXiv Search", "Chat with Claude", "History"])
+
+    with arxiv_tab:
+        st.header("🔍 ArXiv Search")
+        arxiv_query = st.text_area("Enter your research query:", height=100)
+        if st.button("Search ArXiv"):
+            if arxiv_query:
+                with st.spinner("Searching ArXiv..."):
+                    result_mixtral, result_mistral, result_rag = search_arxiv(arxiv_query)
+                    
+                    if result_mixtral:
+                        st.subheader("Mixtral Model Response")
+                        st.markdown(result_mixtral)
+                        
+                        st.subheader("Mistral Model Response")
+                        st.markdown(result_mistral)
+                        
+                        st.subheader("RAG-Enhanced Response")
+                        if isinstance(result_rag, (list, tuple)) and len(result_rag) > 0:
+                            st.markdown(result_rag[0])
+                            if len(result_rag) > 1:
+                                st.markdown(result_rag[1])
+                        
+                        # Save results
+                        filename = generate_filename(arxiv_query, "md")
+                        create_file(filename, arxiv_query, f"{result_mixtral}\n\n{result_mistral}")
+                        
+                        if container:
+                            save_to_cosmos_db(container, arxiv_query, result_mixtral, result_mistral)
+
+    with claude_tab:
+        st.header("💬 Chat with Claude")
+        user_input = st.text_area("Your message:", height=100)
+        if st.button("Send"):
+            if user_input:
+                with st.spinner("Claude is thinking..."):
+                    try:
+                        response = anthropic_client.messages.create(
+                            model="claude-3-sonnet-20240229",
+                            max_tokens=1000,
+                            messages=[{"role": "user", "content": user_input}]
+                        )
+                        
+                        claude_response = response.content[0].text
+                        st.markdown("### Claude's Response:")
+                        st.markdown(claude_response)
+                        
+                        # Save chat history
+                        st.session_state.chat_history.append({
+                            "user": user_input,
+                            "claude": claude_response,
+                            "timestamp": datetime.now().isoformat()
+                        })
+                        
+                        # Save to file
+                        filename = generate_filename(user_input, "md")
+                        create_file(filename, user_input, claude_response)
+                        
+                        # Save to Cosmos DB
+                        if container:
+                            save_to_cosmos_db(container, user_input, claude_response, "")
+                            
+                    except Exception as e:
+                        st.error(f"Error communicating with Claude: {str(e)}")
+
+    with history_tab:
+        st.header("📜 Chat History")
+        for chat in reversed(st.session_state.chat_history):
+            with st.expander(f"Conversation from {chat.get('timestamp', 'Unknown time')}"):
+                st.markdown("**Your message:**")
+                st.markdown(chat["user"])
+                st.markdown("**Claude's response:**")
+                st.markdown(chat["claude"])
+
+if __name__ == "__main__":
+    main()