Prat0 commited on
Commit
2e3951e
·
verified ·
1 Parent(s): ede134b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +136 -0
app.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import os
3
+ from llama_index.core.indices.vector_store.base import VectorStoreIndex
4
+ from llama_index.vector_stores.qdrant import QdrantVectorStore
5
+ from llama_index.embeddings.fastembed import FastEmbedEmbedding
6
+ from langchain_google_genai import ChatGoogleGenerativeAI
7
+ from llama_index.core import Settings
8
+ from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
9
+ import qdrant_client
10
+ from llama_index.core.indices.query.schema import QueryBundle
11
+ from llama_index.llms.gemini import Gemini
12
+ from llama_index.embeddings.gemini import GeminiEmbedding
13
+ from llama_index.core.memory import ChatMemoryBuffer
14
+ from llama_index.readers.web import FireCrawlWebReader
15
+ from llama_index.core import SummaryIndex
16
+
17
+ # Setup functions
18
+ def embed_setup():
19
+ Settings.embed_model = GeminiEmbedding(api_key=os.getenv("GOOGLE_API_KEY"), model_name="models/embedding-001")
20
+ Settings.llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.1, model_name="models/gemini-pro")
21
+
22
+ def qdrant_setup():
23
+ client = qdrant_client.QdrantClient(
24
+ os.getenv('QDRANT_URL'),
25
+ api_key = os.getenv('QDRANT_API_KEY'),
26
+ )
27
+ return client
28
+
29
+ def llm_setup():
30
+ llm = Gemini(api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.1, model_name="models/gemini-pro")
31
+ return llm
32
+
33
+ def query_index(index, similarity_top_k=3, streaming=True):
34
+ memory = ChatMemoryBuffer.from_defaults(token_limit=4000)
35
+ chat_engine = index.as_chat_engine(
36
+ chat_mode="context",
37
+ memory=memory,
38
+ system_prompt=(
39
+ """You are an AI assistant for developers, specializing in technical documentation. Your task is to provide accurate, concise, and helpful responses based on the given documentation context.
40
+ Context information is below:
41
+ {context_str}
42
+ Always answer based on the information in the context and general knowledge and be precise
43
+ Given this context, please respond to the following user query:
44
+ {query_str}
45
+ Your response should:
46
+
47
+ Directly address the query using information from the context
48
+ Include relevant code examples or direct quotes if applicable
49
+ Mention specific sections or pages of the documentation
50
+ Highlight any best practices or potential pitfalls related to the query
51
+
52
+ After your response, suggest 3 follow-up questions based on the context that the user might find helpful for deeper understanding.
53
+ Your response:"""
54
+ ),
55
+ )
56
+ return chat_engine
57
+
58
+ # Document ingestion function
59
+ def ingest_documents(url):
60
+ firecrawl_reader = FireCrawlWebReader(
61
+ api_key=os.getenv("FIRECRAWL_API_KEY"),
62
+ mode="crawl",
63
+ )
64
+ documents = firecrawl_reader.load_data(url=url)
65
+ return documents
66
+
67
+ # Streamlit app
68
+ st.title("Talk to Software Documentation")
69
+
70
+ # Initialize session state
71
+ if 'chat_engine' not in st.session_state:
72
+ st.session_state['chat_engine'] = None
73
+ if 'documents' not in st.session_state:
74
+ st.session_state['documents'] = None
75
+ if 'chat_history' not in st.session_state:
76
+ st.session_state['chat_history'] = []
77
+ if 'last_response' not in st.session_state:
78
+ st.session_state['last_response'] = None
79
+
80
+
81
+ # URL input for document ingestion
82
+ url = st.text_input("Enter URL to crawl and ingest documents:")
83
+
84
+ # Ingest documents button
85
+ if st.button("Ingest Documents"):
86
+ if url:
87
+ with st.spinner("Crawling and ingesting documents..."):
88
+ st.session_state['documents'] = ingest_documents(url)
89
+ st.success(f"Documents ingested from {url}")
90
+ else:
91
+ st.error("Please enter a URL")
92
+
93
+ # Setup button
94
+ if st.button("Setup Query Engine"):
95
+ if st.session_state['documents'] is None:
96
+ st.error("Please ingest documents first")
97
+ else:
98
+ with st.spinner("Setting up query engine..."):
99
+ embed_setup()
100
+ client = qdrant_setup()
101
+ llm = llm_setup()
102
+ vector_store = QdrantVectorStore(client=client, collection_name=os.getenv("COLLECTION_NAME"))
103
+ index = VectorStoreIndex.from_documents(st.session_state['documents'], vector_store=vector_store)
104
+ st.session_state['chat_engine'] = query_index(index)
105
+ st.success("Query engine setup completed successfully!")
106
+
107
+ # Query input
108
+ query = st.text_input("Enter your query:")
109
+
110
+ # Search button
111
+ if st.button("Search"):
112
+ if st.session_state['chat_engine'] is None:
113
+ st.error("Please complete the setup first")
114
+ elif query:
115
+ with st.spinner("Searching..."):
116
+ response = st.session_state['chat_engine'].chat(query)
117
+
118
+ # Add the query and response to chat history
119
+ st.session_state['chat_history'].append(("User", query))
120
+ st.session_state['chat_history'].append(("Assistant", str(response.response)))
121
+
122
+ # Display the most recent response prominently
123
+ st.subheader("Assistant's Response:")
124
+ st.write(response.response)
125
+ else:
126
+ st.error("Please enter a query")
127
+
128
+ if st.session_state['chat_history']:
129
+ st.subheader("Chat History")
130
+ for role, message in st.session_state['chat_history']:
131
+ st.text(f"{role}: {message}")
132
+
133
+ # Clear chat history button
134
+ if st.button("Clear Chat History"):
135
+ st.session_state['chat_history'] = []
136
+ st.success("Chat history cleared!")