Create streamlit_app.py

streamlit_app.py

@@ -0,0 +1,204 @@
+import time
+import streamlit as st
+import pandas as pd
+import openai
+import pinecone
+import requests
+import certifi
+from tiktoken import get_encoding
+from pinecone import Pinecone, ServerlessSpec
+from retry import retry
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from requests.exceptions import SSLError
+from urllib3.exceptions import MaxRetryError
+
+# Configuration
+OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
+PINECONE_API_KEY = os.getenv('PINECONE_API_KEY ')
+INDEX_NAME = "gradient-cyber"
+BATCH_SIZE = 100
+
+# Initialize OpenAI
+openai.api_key = OPENAI_API_KEY
+
+# Retry Pinecone initialization if network issues occur
+@retry(tries=5, delay=2, backoff=2)
+def initialize_pinecone():
+    pc = Pinecone(api_key=PINECONE_API_KEY)
+    if INDEX_NAME not in pc.list_indexes().names():
+        pc.create_index(
+            name=INDEX_NAME,
+            dimension=1536,
+            metric='cosine',
+            spec=ServerlessSpec(cloud='aws', region='us-east-1')
+        )
+    return pc.Index(INDEX_NAME)
+
+# Initialize Pinecone with retries
+index = initialize_pinecone()
+
+def generate_embedding(text, OPENAI_API_KEY):
+    headers = {
+        'Authorization': f'Bearer {OPENAI_API_KEY}',
+        'Content-Type': 'application/json',
+    }
+    data = {
+        'model': 'text-embedding-ada-002',
+        'input': text
+    }
+    response = requests.post(
+        'https://api.openai.com/v1/embeddings',
+        headers=headers,
+        json=data,
+        verify=certifi.where()
+    )
+    if response.status_code == 200:
+        result = response.json()
+        tokens_used = result['usage']['total_tokens']
+        return result['data'][0]['embedding'], tokens_used
+    else:
+        st.error(f"Error creating embedding: {response.status_code} - {response.text}")
+        return None, 0
+
+@retry(tries=5, delay=2, backoff=2, exceptions=(SSLError, MaxRetryError))
+def upsert_in_batches(index, vectors):
+    try:
+        index.upsert(vectors=vectors, namespace="ns1")
+    except (SSLError, MaxRetryError) as e:
+        st.error(f"Error upserting batch: {e}")
+        raise
+
+def process_batches(df, batch_size=100):
+    total_tokens_used = 0
+    total_requests = 0
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = []
+        for start in range(0, len(df), batch_size):
+            batch = df.iloc[start:start + batch_size]
+            for i, row in batch.iterrows():
+                text = row['combined_text']
+                future = executor.submit(generate_embedding, text, OPENAI_API_KEY)
+                futures.append((future, start + i, text))
+
+            vectors = []
+            for future, idx, text in futures:
+                embedding, tokens_used = future.result()
+                if embedding is not None:
+                    total_tokens_used += tokens_used
+                    total_requests += 1
+                    vectors.append({'id': f"{idx}", 'values': embedding, 'metadata': {'text': text}})
+
+            if vectors:
+                upsert_future = executor.submit(upsert_in_batches, index, vectors)
+                futures.append(upsert_future)
+                futures = []  # Clear the futures list for the next batch
+    
+    return total_tokens_used, total_requests
+
+def summarize_texts(texts, OPENAI_API_KEY):
+    headers = {
+        'Authorization': f'Bearer {OPENAI_API_KEY}',
+        'Content-Type': 'application/json',
+    }
+    data = {
+        'model': 'text-davinci-003',
+        'prompt': f"Summarize the following texts:\n\n{texts}\n\n",
+        'max_tokens': 150,
+        'temperature': 0.5,
+    }
+    response = requests.post(
+        'https://api.openai.com/v1/completions',
+        headers=headers,
+        json=data,
+        verify=certifi.where()
+    )
+    if response.status_code == 200:
+        result = response.json()
+        return result['choices'][0]['text'].strip()
+    else:
+        st.error(f"Error summarizing text: {response.status_code} - {response.text}")
+        return None
+
+def answer_query(query, summaries, api_key):
+    combined_text = "\n\n".join([f"ID: {summary['hash_id']}\n\n{summary['text']}" for summary in summaries])
+    headers = {
+        'Authorization': f'Bearer {api_key}',
+        'Content-Type': 'application/json',
+    }
+    data = {
+        'model': 'gpt-3.5-turbo',
+        'messages': [
+            {"role": "system",
+             "content": ("Core Capabilities: Expert Knowledge on SITREPs: Understand and explain SITREP components "
+                         "like threat analysis, incident summaries, and risk assessments. Flexible Query Handling: "
+                         "Interpret and respond to diverse queries, supporting access to data by date, theme, severity, etc. "
+                         "Data Retrieval: Provide specific information such as incident dates, threat actors, and mitigation strategies. "
+                         "Offer summaries or detailed reports based on user needs. Analytical Engagement: Engage in discussions, offering insights and hypotheses. "
+                         "Support user analysis of cyber threats and incidents. Interaction Guidelines: Understanding Queries: "
+                         "Use NLP to interpret user queries and ask clarifying questions if needed. Providing Context: Give context to help users understand the relevance and implications of information. "
+                         "Customizing Detail: Adjust detail levels based on user preferences, providing summaries or deep dives. "
+                         "Allow users to specify the format of information (e.g., bullet points, detailed paragraphs, tables). "
+                         "Encouraging Exploration: Suggest related queries and additional information. Provide thorough explanations to support learning and decision-making.")
+            },
+            {"role": "user", "content": f"Here are some summaries \n\n{combined_text}\n\nAnswer this query: {query}"}
+        ],
+        'max_tokens': 1500,
+        'temperature': 0.4,
+    }
+    response = requests.post(
+        'https://api.openai.com/v1/chat/completions',
+        headers=headers,
+        json=data,
+        verify=certifi.where()
+    )
+    if response.status_code == 200:
+        result = response.json()
+        tokens_used = result['usage']['total_tokens']
+        answer = result['choices'][0]['message']['content']
+        used_hash_ids = [summary['hash_id'] for summary in summaries if summary['hash_id'] in result['choices'][0]['message']['content']]
+        return answer, used_hash_ids, tokens_used
+    else:
+        return f"Error: {response.status_code} - {response.text}", [], 0
+
+# Streamlit UI for file upload
+st.title("Excel File Upload and Processing")
+uploaded_file = st.file_uploader("Upload an Excel file", type=["xlsx"])
+if uploaded_file is not None:
+    st.write("File uploaded successfully!")
+    # Load Excel file
+    df = pd.read_excel(uploaded_file)
+    st.write("Excel file loaded:")
+    st.write(df.head())
+    # Concatenate text from all columns for each row
+    def create_meaningful_sentence(row):
+        return '. '.join([f"{col}: {row[col]}" for col in df.columns])
+    df['combined_text'] = df.apply(create_meaningful_sentence, axis=1)
+    st.write("Columns concatenated into meaningful sentences:")
+    st.write(df[['combined_text']].head())
+    
+    # Process rows in batches and upsert vectors
+    total_tokens_used, total_requests = process_batches(df, BATCH_SIZE)
+    
+    st.success(f"Data successfully uploaded to Pinecone.")
+    st.info(f"Total tokens used: {total_tokens_used}")
+    st.info(f"Total requests made: {total_requests}")
+
+# Query input and response
+st.title("Query the Uploaded Data")
+query = st.text_input("Enter your query:")
+if query:
+    query_embedding, _ = generate_embedding(query, OPENAI_API_KEY)
+    if query_embedding is not None:
+        # Search Pinecone with the query embedding
+        results = index.query(namespace="ns1", include_values=True, vector=query_embedding, top_k=5, include_metadata=True)
+        # Collect texts for summarization
+        summaries = [{'hash_id': match['id'], 'text': match['metadata']['text']} for match in results['matches']]
+        answer, used_hash_ids, tokens_used = answer_query(query, summaries, OPENAI_API_KEY)
+        # Display results
+        if answer:
+            st.write("Answer to your query:")
+            st.write(answer)
+            st.info(f"Tokens used for this query: {tokens_used}")
+        else:
+            st.write("No answer could be generated.")