Ransaka/Code-Assistant
0
1import redis 2import os 3import google.generativeai as genai4from typing import List5import numpy as np 6from redis.commands.search.query import Query7from haystack import Pipeline, component8from haystack.utils import Secret9from haystack_integrations.components.generators.google_ai import GoogleAIGeminiChatGenerator, GoogleAIGeminiGenerator10from haystack.components.builders import PromptBuilder11import streamlit as st12from data_processor import fetch_data,ingest_data13 14genai.configure(api_key=os.environ["GEMINI_API_KEY"])15 16generation_config = {17 "temperature": 1,18 "top_p": 0.95,19 "top_k": 64,20 "max_output_tokens": 8192,21 "response_mime_type": "text/plain",22}23 24model = genai.GenerativeModel(25 model_name="gemini-1.5-flash",26 generation_config=generation_config,27 system_instruction="You are optimized to generate accurate descriptions for given Python codes. When the user inputs the code, you must return the description according to its goal and functionality. You are not allowed to generate additional details. The user expects at least 5 sentence-long descriptions.",28)29 30gemini = GoogleAIGeminiGenerator(api_key=Secret.from_env_var("GEMINI_API_KEY"), model='gemini-1.5-flash')31 32def get_embeddings(content: List):33 return genai.embed_content(model='models/text-embedding-004',content=content)['embedding']34 35 36def draft_prompt(query: str, chat_history: str) -> str:37 """38 Perform a vector similarity search and retrieve related functions.39 40 Args:41 query (str): The input query to encode.42 43 Returns:44 str: A formatted string containing details of related functions.45 """46 INDEX_NAME = "idx:codes_vss"47 client = st.session_state.client48 vector_search_query = (49 Query('(*)=>[KNN 2 @vector $query_vector AS vector_score]')50 .sort_by('vector_score')51 .return_fields('vector_score', 'id', 'name', 'definition', 'file_name', 'type', 'uses')52 .dialect(2)53 )54 55 encoded_query = get_embeddings(query)56 vector_params = {57 "query_vector": np.array(encoded_query, dtype=np.float32).tobytes()58 }59 60 result_docs = client.ft(INDEX_NAME).search(vector_search_query, vector_params).docs61 62 related_items: List[str] = []63 dependencies: List[str] = []64 for doc in result_docs:65 related_items.append(doc.name)66 if doc.uses:67 dependencies.extend(use for use in doc.uses.split(", ") if use)68 69 dependencies = list(set(dependencies) - set(related_items))70 71 def get_query(item_list):72 return Query(f"@name:({' | '.join(item_list)})").return_fields(73 'id', 'name', 'definition', 'file_name', 'type'74 )75 76 related_docs = client.ft(INDEX_NAME).search(get_query(related_items)).docs77 dependency_docs = client.ft(INDEX_NAME).search(get_query(dependencies)).docs78 79 def format_doc(doc):80 return (81 f"{'*' * 28} CODE SNIPPET {doc.id} {'*' * 28}\n"82 f"* Name: {doc.name}\n"83 f"* File: {doc.file_name}\n"84 f"* {doc.type.capitalize()} definition:\n"85 f"```python\n{doc.definition}\n```\n"86 )87 88 formatted_results_main = [format_doc(doc) for doc in related_docs]89 formatted_results_support = [format_doc(doc) for doc in dependency_docs]90 91 return (92 f"User Question: {query}\n\n"93 f"Current Chat History: \n{chat_history}\n\n"94 f"USE BELOW CODES TO ANSWER USER QUESTIONS.\n"95 f"{chr(10).join(formatted_results_main)}\n\n"96 f"SOME SUPPORTING FUNCTIONS AND CLASS YOU MAY WANT.\n"97 f"{chr(10).join(formatted_results_support)}"98 )99 100@component101class RedisRetreiver:102 @component.output_types(context=str)103 def run(self, query:str, chat_history:str):104 return {"context": draft_prompt(query, chat_history)}105 106llm = GoogleAIGeminiGenerator(api_key=Secret.from_env_var("GEMINI_API_KEY"), model='gemini-1.5-pro')107# llm = OpenAIGenerator()108 109template = """110You are a helpful agent optimized to resolve GitHub issues for your organization's libraries. Users will ask questions when they encounter problems with the code repository.111You have access to all the necessary code for addressing these issues. 112First, you should understand the user's question and identify the relevant code blocks. 113Then, craft a precise and targeted response that allows the user to find an exact solution to their problem. 114You must provide code snippets rather than just opinions.115You should always assume user has installed this python package in their system and raised question raised while they are using the library.116 117In addition to the above tasks, you are free to:118 * Greet the user.119 * [ONLY IF THE QUESTION IS INSUFFICIENT] Request additional clarity.120 * Politely decline irrelevant queries.121 * Inform the user if their query cannot be processed or accomplished.122 123By any chance you should NOT,124 * Ask or recommend user to use different library. Or code snipits related to other similar libraies.125 * Provide inaccurate explnations.126 * Provide sugestions without code examples.127 128{{context}}129"""130 131prompt_builder = PromptBuilder(template=template)132 133pipeline = Pipeline()134pipeline.add_component(name="retriever", instance=RedisRetreiver())135pipeline.add_component("prompt_builder", prompt_builder)136pipeline.add_component("llm", llm)137pipeline.connect("retriever.context", "prompt_builder")138pipeline.connect("prompt_builder", "llm")139 140# Initialize Streamlit app141st.title("Code Assistant Chat")142st.subheader("Frequently Asked Questions")143 144st.markdown("""145 <style>146 .streamlit-expanderHeader {147 background-color: #f0f2f6;148 border: 1px solid #ddd;149 border-radius: 5px;150 padding: 10px;151 }152 .streamlit-expanderContent {153 background-color: #ffffff;154 border: 1px solid #ddd;155 border-radius: 5px;156 padding: 10px;157 }158 </style>159""", unsafe_allow_html=True)160 161with st.expander("How can I use this space?"):162 st.markdown("""163 This space is created based on steps described in [this Medium article](https://towardsdatascience.com/building-llm-powered-coding-assitant-for-github-b88beeb42f2d). To use this space:164 165 1. Create a Redis Cloud account and set up a database166 2. Add your Redis credentials to this space167 3. Enter your preferred GitHub repository clone URL for data fetching and indexing168 169 """)170 171with st.expander("Do I need a Gemini API key?"):172 st.markdown("""173 No, you don't need to provide a Gemini API key for testing this repository. 174 175 - This repo includes a Gemini free tier API key itself. 176 - However, if you encounter any resource exhaustion error:177 - Consider cloning this space178 - Add your own key as the `GEMINI_API_KEY` secret179 180 """)181 182with st.expander("I don't want to create a Redis database. Can I still check the output?"):183 st.markdown("""184 Absolutely! Here's what you can do:185 186 1. Send me a message on [LinkedIn](https://www.linkedin.com/in/ransaka/) mentioning your requirement187 2. I'll provide you with preconfigured database credentials188 3. Enter these credentials in the appropriate fields189 4. You'll then be able to use the assistant as you wish190 191 > **Important**: Please use the provided credentials responsibly and for testing purposes only.192 """)193 194tabs = ["Data Fetching","Assistant"]195selected_tab = st.sidebar.radio("Select a Tab", tabs)196if selected_tab == 'Data Fetching':197 if 'redis_connected' not in st.session_state:198 st.session_state.redis_connected = False199 200 if not st.session_state.redis_connected:201 st.header("Redis Connection Settings")202 203 redis_username = st.text_input("Redis Username", value='default')204 redis_host = st.text_input("Redis Host")205 redis_port = st.number_input("Redis Port", min_value=1, max_value=65535, value=5555)206 redis_password = st.text_input("Redis Password", type="password")207 208 if st.button("Connect to Redis"):209 try:210 client = redis.Redis(211 host=redis_host,212 port=redis_port,213 password=redis_password,214 username=redis_username215 )216 217 if client.ping():218 st.success("Successfully connected to Redis!")219 st.session_state.redis_connected = True220 st.session_state.client = client221 st.session_state.host = redis_host222 else:223 st.error("Failed to connect to Redis. Please check your settings.")224 except redis.ConnectionError:225 st.error("Failed to connect to Redis. Please check your settings and try again.")226 227 if st.session_state.redis_connected:228 if st.session_state.host == os.environ['REDIS_HOST']:229 st.success("You are all set!")230 else:231 url = st.text_input("Enter git clone URL")232 if url:233 with st.spinner("Fetching data..."):234 data = fetch_data(url)235 236 with st.spinner("Ingesting data..."):237 response_string = ingest_data(st.session_state.client, data)238 239 st.write(response_string)240 241if selected_tab == 'Assistant':242 if "messages" not in st.session_state:243 st.session_state.messages = []244 245 # Display chat messages246 for message in st.session_state.messages:247 with st.chat_message(message["role"]):248 st.markdown(message["content"])249 250 st.session_state.response = None251 252 if prompt := st.chat_input("What's your question?"):253 st.session_state.messages.append({"role": "user", "content": prompt})254 with st.chat_message("user"):255 st.markdown(prompt)256 257 with st.chat_message("assistant"):258 response_placeholder = st.empty()259 response_placeholder.markdown("Thinking...")260 261 try:262 response = pipeline.run({"retriever": {"query": prompt, "chat_history": st.session_state.messages}}, include_outputs_from=['prompt_builder'])263 st.session_state.response = response264 llm_response = response["llm"]["replies"][0]265 266 response_placeholder.markdown(llm_response)267 st.session_state.messages.append({"role": "assistant", "content": llm_response})268 except Exception as e:269 response_placeholder.markdown(f"An error occurred: {str(e)}")270 271 if st.button("Clear Chat History"):272 st.session_state.messages = []273 st.experimental_rerun()274 275 with st.expander("See Chat History"):276 st.markdown(st.session_state.response)