import streamlit as st import pandas as pd from openai import OpenAI import requests import certifi import tiktoken from tiktoken import get_encoding from pinecone import Pinecone, ServerlessSpec import time # Configuration import os from dotenv import load_dotenv # Load environment variables from .env file load_dotenv() # Configuration OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') PINECONE_API_KEY = "9c097a58-6008-409a-859a-668a002320f6" INDEX_NAME = "gradient-cyber" BATCH_SIZE = 100 MAX_RESULTS = 1000 # Initialize OpenAI client = OpenAI(api_key=OPENAI_API_KEY) # Initialize Pinecone pc = Pinecone(api_key=PINECONE_API_KEY) # Check if the index already exists before creating it 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') ) index = pc.Index(INDEX_NAME) # Define helper functions def truncate_text(text, max_tokens): tokenizer = get_encoding("gpt2") tokens = tokenizer.encode(text) return tokenizer.decode(tokens[:max_tokens]) def generate_embedding(text): max_retries = 3 for attempt in range(max_retries): try: response = client.embeddings.create( model="text-embedding-ada-002", input=text ) return response.data[0].embedding, response.usage.total_tokens except Exception as e: if attempt == max_retries - 1: st.error(f"Error creating embedding after {max_retries} attempts: {str(e)}") return None, 0 time.sleep(2 ** attempt) # Exponential backoff def upsert_in_batches(index, vectors, batch_size=100): batches = [vectors[i:i + batch_size] for i in range(0, len(vectors), batch_size)] for batch in batches: try: index.upsert(vectors=batch, namespace="ns1") except Exception as e: st.error(f"Error upserting batch: {e}") def num_tokens_from_string(string: str, encoding_name: str = "cl100k_base") -> int: encoding = tiktoken.get_encoding(encoding_name) num_tokens = len(encoding.encode(string)) return num_tokens def semantic_similarity(text1, text2): embedding1, _ = generate_embedding(text1) embedding2, _ = generate_embedding(text2) if embedding1 is None or embedding2 is None: return 0 return sum(a*b for a, b in zip(embedding1, embedding2)) def expand_query(original_query): try: expansion_prompt = f"Expand the following query into 3-5 related questions or terms: '{original_query}'" response = client.chat.completions.create( model="gpt-4o-mini", messages=[{"role": "user", "content": expansion_prompt}], max_tokens=100, temperature=0.7 ) return original_query + " " + response.choices[0].message.content except Exception as e: st.error(f"Error in query expansion: {str(e)}") return original_query # Return original query if expansion fails def truncate_context(context, max_tokens=14000): encoding = tiktoken.get_encoding("cl100k_base") encoded = encoding.encode(context) truncated = encoded[:max_tokens] return encoding.decode(truncated) # Streamlit UI for file upload st.title("Gradient-cyber") 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 into a readable sentence def create_meaningful_sentence(row): return '. '.join([f"{col.replace('_', ' ')}: {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()) vectors = [] # Process each row in the DataFrame total_tokens_used = 0 total_requests = 0 for i, row in df.iterrows(): text = row['combined_text'] # Truncate text to fit within the model's maximum context length text = truncate_text(text, max_tokens=8192) embedding, tokens_used = generate_embedding(text) if embedding is not None: total_tokens_used += tokens_used total_requests += 1 # Truncate text fields to reduce metadata size def truncate_field(field, max_length=500): return str(field)[:max_length] if not pd.isna(field) else '' # Prepare metadata with handling NaN values and converting to string metadata = { "ID": truncate_field(row['ID']), "eventDtgTime": truncate_field(row['eventDtgTime']), "alerts": truncate_field(row.get('alerts', '')), "displayTitle": truncate_field(row['displayTitle']), "instantAnalytics": truncate_field(row.get('instantAnalytics', '')), "detailedText": truncate_field(row.get('detailedText', '')), "msgPrecs": truncate_field(row.get('msgPrecs', '')), "unit": truncate_field(row.get('unit', '')), "size": truncate_field(row.get('size', '')), "embedHtml": truncate_field(row.get('embedHtml', '')), "dataSources": truncate_field(row.get('dataSources', '')), "snippetText": truncate_field(row.get('snippetText', '')), "contentLink": truncate_field(row.get('contentLink', '')), "description": truncate_field(row.get('description', '')), "imageDescription": truncate_field(row.get('imageDescription', '')), "reportSummary": truncate_field(row.get('reportSummary', '')), "authorName": truncate_field(row.get('authorName', '')), "timeReportCompleted": truncate_field(row.get('timeReportCompleted', '')), "attachment": truncate_field(row.get('attachment', '')), "latitude": truncate_field(row.get('latitude', '')), "securityLevels": truncate_field(row.get('securityLevels', '')), "imagereSourceLink": truncate_field(row.get('imagereSourceLink', '')), "eventDtg": truncate_field(row.get('eventDtg', '')), "status": truncate_field(row.get('status', '')), "users": truncate_field(row.get('users', '')), "name": truncate_field(row.get('name', '')), "sessions": truncate_field(row.get('sessions', '')), "fiscalStatus": truncate_field(row.get('fiscalStatus', '')), "sentimentSummary": truncate_field(row.get('sentimentSummary', '')), "sourceOrg": truncate_field(row.get('sourceOrg', '')), "dateCreated": truncate_field(row.get('dateCreated', '')), "active": truncate_field(row.get('active', '')), "responseSummary": truncate_field(row.get('responseSummary', '')), "comparisonCommunitiesCountries": truncate_field(row.get('comparisonCommunitiesCountries', '')), "activity": truncate_field(row.get('activity', '')), "applications": truncate_field(row.get('applications', '')), "url": truncate_field(row.get('url', '')), "timeZones": truncate_field(row.get('timeZones', '')), "location": truncate_field(row.get('location', '')), "longitude": truncate_field(row.get('longitude', '')), "dateModified": truncate_field(row.get('dateModified', '')), "pedigrees": truncate_field(row.get('pedigrees', '')), "gistComment": truncate_field(row.get('gistComment', '')), "tag": truncate_field(row.get('tag', '')), "geoCode": truncate_field(row.get('geoCode', '')), "time": truncate_field(row.get('time', '')), "timeReportRouted": truncate_field(row.get('timeReportRouted', '')), "rteToOrg": truncate_field(row.get('rteToOrg', '')), "copyReportToOrg": truncate_field(row.get('copyReportToOrg', '')), "sourceOrganization": truncate_field(row.get('sourceOrganization', '')), "coordinates": truncate_field(row.get('coordinates', '')), "image1": truncate_field(row.get('image1', '')), "image2": truncate_field(row.get('image2', '')), "image3": truncate_field(row.get('image3', '')), "image4": truncate_field(row.get('image4', '')), "image5": truncate_field(row.get('image5', '')), "numEmailsSent": truncate_field(row.get('numEmailsSent', '')), "lastEmailDate": truncate_field(row.get('lastEmailDate', '')), "reportDtg": truncate_field(row.get('reportDtg', '')), "metadata": truncate_field(row.get('metadata', '')), "eventOrganizations": truncate_field(row.get('eventOrganizations', '')), "classification": truncate_field(row.get('classification', '')), "assetIPs": truncate_field(row.get('assetIPs', '')), "sitrepTemplate": truncate_field(row.get('sitrepTemplate', '')), "industry": truncate_field(row.get('industry', '')), "networkSegmentList": truncate_field(row.get('networkSegmentList', '')), "approvedDate": truncate_field(row.get('approvedDate', '')), "incident": truncate_field(row.get('incident', '')), "sendEmail": truncate_field(row.get('sendEmail', '')), "newFormat": truncate_field(row.get('newFormat', '')), "duMapping": truncate_field(row.get('duMapping', '')), "jsonTag": truncate_field(row.get('jsonTag', '')), "createdFrom": truncate_field(row.get('createdFrom', '')), "integrationData": truncate_field(row.get('integrationData', '')), "mtti": truncate_field(row.get('mtti', '')), "mttd": truncate_field(row.get('mttd', '')), "mttr": truncate_field(row.get('mttr', '')), "oldEventDate": truncate_field(row.get('oldEventDate', '')), "org_event_name": truncate_field(row.get('org_event_name', '')), "combined_text": text # Add combined text to metadata } vectors.append({'id': str(row['ID']), 'values': embedding, 'metadata': metadata}) if vectors: upsert_in_batches(index, vectors, 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}") else: st.warning("No embeddings were generated.") # Query input and response query = st.text_input("Enter your query:") if query: try: expanded_query = expand_query(query) query_embedding, _ = generate_embedding(expanded_query) if query_embedding is not None: # Perform the search in Pinecone results = index.query( namespace="ns1", vector=query_embedding, top_k=50, include_metadata=True ) # Semantic filtering filtered_results = sorted( results['matches'], key=lambda x: semantic_similarity(query, x['metadata']['combined_text']), reverse=True )[:10] # Reduced from 20 to 15 to further limit context size # Prepare context for GPT context = "\n".join([ f"ID: {match['id']}\n" + f"Event Date/Time: {match['metadata'].get('eventDtgTime', 'N/A')}\n" + f"Display Title: {match['metadata'].get('displayTitle', 'N/A')}\n" + f"Status: {match['metadata'].get('status', 'N/A')}\n" + f"Combined Text: {match['metadata'].get('combined_text', 'N/A')}\n" + "---" for match in filtered_results ]) # Truncate the context truncated_context = truncate_context(context) # Prepare the prompt for GPT system_prompt = """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.""" user_prompt = f"""Query: {query} Relevant Information: {truncated_context} Provide a clear, concise, and comprehensive answer. Synthesize information from multiple entries if necessary. Cite specific details and examples when applicable. If information is missing, state what is known and what remains uncertain.""" response = client.chat.completions.create( model="gpt-4o-mini", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ], max_tokens=1000, temperature=0.7 ) answer = response.choices[0].message.content st.write("Answer to your query:") st.write(answer) else: st.error("Failed to generate query embedding.") except Exception as e: st.error(f"Error processing query: {str(e)}")