code: # Copyright (c) Timescale, Inc. (2023)
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import time
  
  import numpy as np
  
  import streamlit as st
  from streamlit.hello.utils import show_code
  
  from llama_index.vector_stores import TimescaleVectorStore
  from llama_index import ServiceContext, StorageContext
  from llama_index.indices.vector_store import VectorStoreIndex
  from llama_index.llms import OpenAI
  from llama_index import set_global_service_context
  
  import pandas as pd
  from pathlib import Path
  from datetime import datetime, timedelta
  from timescale_vector import client
  
  from typing import List, Tuple
  
  from llama_index.schema import TextNode
  from llama_index.embeddings import OpenAIEmbedding
  import psycopg2
  
  def get_repos():
      with psycopg2.connect(dsn=st.secrets["TIMESCALE_SERVICE_URL"]) as connection:
          # Create a cursor within the context manager
          with connection.cursor() as cursor:
              try:
                  select_data_sql = "SELECT * FROM time_machine_catalog;"
                  cursor.execute(select_data_sql)
              except psycopg2.errors.UndefinedTable as e:
                  return {}
  
              catalog_entries = cursor.fetchall()
  
              catalog_dict = {}
              for entry in catalog_entries:
                  repo_url, table_name = entry
                  catalog_dict[repo_url] = table_name
  
              return catalog_dict
  
  def get_auto_retriever(index, retriever_args):
      from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo
      vector_store_info = VectorStoreInfo(
          content_info="Description of the commits to PostgreSQL. Describes changes made to Postgres",
          metadata_info=[
              MetadataInfo(
                  name="commit_hash",
                  type="str",
                  description="Commit Hash",
              ),
              MetadataInfo(
                  name="author",
                  type="str",
                  description="Author of the commit",
              ),
              MetadataInfo(
                  name="__start_date",
                  type="datetime in iso format",
                  description="All results will be after this datetime",
      
              ),
              MetadataInfo(
                  name="__end_date",
                  type="datetime in iso format",
                  description="All results will be before this datetime",
      
              )
          ],
      )
      from llama_index.indices.vector_store.retrievers import VectorIndexAutoRetriever
      retriever = VectorIndexAutoRetriever(index, 
                                           vector_store_info=vector_store_info, 
                                           service_context=index.service_context,
                                           **retriever_args)
      
      # build query engine
      from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine
      query_engine = RetrieverQueryEngine.from_args(
          retriever=retriever, service_context=index.service_context
      )
  
      from llama_index.tools.query_engine import QueryEngineTool
      # convert query engine to tool
      query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
  
      from llama_index.agent import OpenAIAgent
      chat_engine = OpenAIAgent.from_tools(
          tools=[query_engine_tool],
          llm=index.service_context.llm,
          verbose=True
          #service_context=index.service_context
      )
      return chat_engine
  
  def tm_demo():
      repos = get_repos()
  
      months = st.sidebar.slider('How many months back to search (0=no limit)?', 0, 130, 0)
  
      if "config_months" not in st.session_state.keys() or months != st.session_state.config_months:
          st.session_state.clear()
  
      topk = st.sidebar.slider('How many commits to retrieve', 1, 150, 20)
      if "config_topk" not in st.session_state.keys() or topk != st.session_state.config_topk:
          st.session_state.clear()
          
      if len(repos) > 0:
          repo = st.sidebar.selectbox("Choose a repo", repos.keys())
      else:
          st.error("No repositiories found, please [load some data first](/LoadData)")
          return
      
      if "config_repo" not in st.session_state.keys() or repo != st.session_state.config_repo:
          st.session_state.clear()
      
      st.session_state.config_months = months
      st.session_state.config_topk = topk
      st.session_state.config_repo = repo
  
  
      if "messages" not in st.session_state.keys(): # Initialize the chat messages history
          st.session_state.messages = [
              {"role": "assistant", "content": "Please choose a repo and time filter on the sidebar and then ask me a question about the git history"}
          ]
  
      vector_store = TimescaleVectorStore.from_params(
          service_url=st.secrets["TIMESCALE_SERVICE_URL"],
          table_name=repos[repo],
          time_partition_interval=timedelta(days=7),
      );
  
      service_context = ServiceContext.from_defaults(llm=OpenAI(model="gpt-4", temperature=0.1))
      set_global_service_context(service_context)
      index = VectorStoreIndex.from_vector_store(vector_store=vector_store, service_context=service_context)
      
          
      #chat engine goes into the session to retain history
      if "chat_engine" not in st.session_state.keys(): # Initialize the chat engine
          retriever_args = {"similarity_top_k" : int(topk)}
          if months > 0:
              end_dt = datetime.now()
              start_dt = end_dt - timedelta(weeks=4*months)
              retriever_args["vector_store_kwargs"] = ({"start_date": start_dt, "end_date":end_dt})
          st.session_state.chat_engine = get_auto_retriever(index, retriever_args)
          #st.session_state.chat_engine = index.as_chat_engine(chat_mode="best", similarity_top_k=20, verbose=True)
  
      if prompt := st.chat_input("Your question"): # Prompt for user input and save to chat history
          st.session_state.messages.append({"role": "user", "content": prompt})
  
      for message in st.session_state.messages: # Display the prior chat messages
          with st.chat_message(message["role"]):
              st.write(message["content"])
  
      # If last message is not from assistant, generate a new response
      if st.session_state.messages[-1]["role"] != "assistant":
          with st.chat_message("assistant"):
              with st.spinner("Thinking..."):
                  response = st.session_state.chat_engine.chat(prompt, function_call="query_engine_tool")
                  st.write(response.response)
                  message = {"role": "assistant", "content": response.response}
                  st.session_state.messages.append(message) # Add response to message history
  
  st.set_page_config(page_title="Time machine demo", page_icon="🧑‍💼")
  st.markdown("# Time Machine")
  st.sidebar.header("Welcome to the Time Machine")
  
  debug_llamaindex = False
  if debug_llamaindex:
      import logging
      import sys
  
      logging.basicConfig(stream=sys.stdout, level=logging.INFO)
      logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
  
  tm_demo()
  
  #show_code(tm_demo)
  
apis: llama_index.tools.query_engine.QueryEngineTool.from_defaults llama_index.indices.vector_store.VectorStoreIndex.from_vector_store llama_index.indices.vector_store.retrievers.VectorIndexAutoRetriever llama_index.llms.OpenAI llama_index.vector_stores.types.MetadataInfo llama_index.agent.OpenAIAgent.from_tools llama_index.query_engine.retriever_query_engine.RetrieverQueryEngine.from_args llama_index.set_global_service_context
extract_api: [((7098, 7170), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""Time machine demo"""', 'page_icon': '"""🧑\u200d💼"""'}), "(page_title='Time machine demo', page_icon='🧑\\u200d💼')\n", (7116, 7170), True, 'import streamlit as st\n'), ((7166, 7195), 'streamlit.markdown', 'st.markdown', (['"""# Time Machine"""'], {}), "('# Time Machine')\n", (7177, 7195), True, 'import streamlit as st\n'), ((7196, 7244), 'streamlit.sidebar.header', 'st.sidebar.header', (['"""Welcome to the Time Machine"""'], {}), "('Welcome to the Time Machine')\n", (7213, 7244), True, 'import streamlit as st\n'), ((2991, 3120), 'llama_index.indices.vector_store.retrievers.VectorIndexAutoRetriever', 'VectorIndexAutoRetriever', (['index'], {'vector_store_info': 'vector_store_info', 'service_context': 'index.service_context'}), '(index, vector_store_info=vector_store_info,\n    service_context=index.service_context, **retriever_args)\n', (3015, 3120), False, 'from llama_index.indices.vector_store.retrievers import VectorIndexAutoRetriever\n'), ((3376, 3471), 'llama_index.query_engine.retriever_query_engine.RetrieverQueryEngine.from_args', 'RetrieverQueryEngine.from_args', ([], {'retriever': 'retriever', 'service_context': 'index.service_context'}), '(retriever=retriever, service_context=index.\n    service_context)\n', (3406, 3471), False, 'from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine\n'), ((3604, 3660), 'llama_index.tools.query_engine.QueryEngineTool.from_defaults', 'QueryEngineTool.from_defaults', ([], {'query_engine': 'query_engine'}), '(query_engine=query_engine)\n', (3633, 3660), False, 'from llama_index.tools.query_engine import QueryEngineTool\n'), ((3726, 3825), 'llama_index.agent.OpenAIAgent.from_tools', 'OpenAIAgent.from_tools', ([], {'tools': '[query_engine_tool]', 'llm': 'index.service_context.llm', 'verbose': '(True)'}), '(tools=[query_engine_tool], llm=index.service_context\n    .llm, verbose=True)\n', (3748, 3825), False, 'from llama_index.agent import OpenAIAgent\n'), ((3975, 4051), 'streamlit.sidebar.slider', 'st.sidebar.slider', (['"""How many months back to search (0=no limit)?"""', '(0)', '(130)', '(0)'], {}), "('How many months back to search (0=no limit)?', 0, 130, 0)\n", (3992, 4051), True, 'import streamlit as st\n'), ((4197, 4258), 'streamlit.sidebar.slider', 'st.sidebar.slider', (['"""How many commits to retrieve"""', '(1)', '(150)', '(20)'], {}), "('How many commits to retrieve', 1, 150, 20)\n", (4214, 4258), True, 'import streamlit as st\n'), ((5443, 5486), 'llama_index.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (5469, 5486), False, 'from llama_index import set_global_service_context\n'), ((5499, 5597), 'llama_index.indices.vector_store.VectorStoreIndex.from_vector_store', 'VectorStoreIndex.from_vector_store', ([], {'vector_store': 'vector_store', 'service_context': 'service_context'}), '(vector_store=vector_store,\n    service_context=service_context)\n', (5533, 5597), False, 'from llama_index.indices.vector_store import VectorStoreIndex\n'), ((7331, 7389), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.INFO'}), '(stream=sys.stdout, level=logging.INFO)\n', (7350, 7389), False, 'import logging\n'), ((1240, 1297), 'psycopg2.connect', 'psycopg2.connect', ([], {'dsn': "st.secrets['TIMESCALE_SERVICE_URL']"}), "(dsn=st.secrets['TIMESCALE_SERVICE_URL'])\n", (1256, 1297), False, 'import psycopg2\n'), ((4160, 4184), 'streamlit.session_state.clear', 'st.session_state.clear', ([], {}), '()\n', (4182, 4184), True, 'import streamlit as st\n'), ((4360, 4384), 'streamlit.session_state.clear', 'st.session_state.clear', ([], {}), '()\n', (4382, 4384), True, 'import streamlit as st\n'), ((4502, 4578), 'streamlit.error', 'st.error', (['"""No repositiories found, please [load some data first](/LoadData)"""'], {}), "('No repositiories found, please [load some data first](/LoadData)')\n", (4510, 4578), True, 'import streamlit as st\n'), ((4700, 4724), 'streamlit.session_state.clear', 'st.session_state.clear', ([], {}), '()\n', (4722, 4724), True, 'import streamlit as st\n'), ((4881, 4904), 'streamlit.session_state.keys', 'st.session_state.keys', ([], {}), '()\n', (4902, 4904), True, 'import streamlit as st\n'), ((5693, 5716), 'streamlit.session_state.keys', 'st.session_state.keys', ([], {}), '()\n', (5714, 5716), True, 'import streamlit as st\n'), ((6233, 6263), 'streamlit.chat_input', 'st.chat_input', (['"""Your question"""'], {}), "('Your question')\n", (6246, 6263), True, 'import streamlit as st\n'), ((6322, 6391), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'user', 'content': prompt}"], {}), "({'role': 'user', 'content': prompt})\n", (6354, 6391), True, 'import streamlit as st\n'), ((7425, 7465), 'logging.StreamHandler', 'logging.StreamHandler', ([], {'stream': 'sys.stdout'}), '(stream=sys.stdout)\n', (7446, 7465), False, 'import logging\n'), ((4083, 4106), 'streamlit.session_state.keys', 'st.session_state.keys', ([], {}), '()\n', (4104, 4106), True, 'import streamlit as st\n'), ((4287, 4310), 'streamlit.session_state.keys', 'st.session_state.keys', ([], {}), '()\n', (4308, 4310), True, 'import streamlit as st\n'), ((4627, 4650), 'streamlit.session_state.keys', 'st.session_state.keys', ([], {}), '()\n', (4648, 4650), True, 'import streamlit as st\n'), ((5317, 5334), 'datetime.timedelta', 'timedelta', ([], {'days': '(7)'}), '(days=7)\n', (5326, 5334), False, 'from datetime import datetime, timedelta\n'), ((5399, 5437), 'llama_index.llms.OpenAI', 'OpenAI', ([], {'model': '"""gpt-4"""', 'temperature': '(0.1)'}), "(model='gpt-4', temperature=0.1)\n", (5405, 5437), False, 'from llama_index.llms import OpenAI\n'), ((5849, 5863), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (5861, 5863), False, 'from datetime import datetime, timedelta\n'), ((6486, 6518), 'streamlit.chat_message', 'st.chat_message', (["message['role']"], {}), "(message['role'])\n", (6501, 6518), True, 'import streamlit as st\n'), ((6532, 6560), 'streamlit.write', 'st.write', (["message['content']"], {}), "(message['content'])\n", (6540, 6560), True, 'import streamlit as st\n'), ((6705, 6733), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (6720, 6733), True, 'import streamlit as st\n'), ((7394, 7413), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (7411, 7413), False, 'import logging\n'), ((2185, 2256), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""commit_hash"""', 'type': '"""str"""', 'description': '"""Commit Hash"""'}), "(name='commit_hash', type='str', description='Commit Hash')\n", (2197, 2256), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((2333, 2408), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""author"""', 'type': '"""str"""', 'description': '"""Author of the commit"""'}), "(name='author', type='str', description='Author of the commit')\n", (2345, 2408), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((2485, 2608), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""__start_date"""', 'type': '"""datetime in iso format"""', 'description': '"""All results will be after this datetime"""'}), "(name='__start_date', type='datetime in iso format',\n    description='All results will be after this datetime')\n", (2497, 2608), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((2686, 2809), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""__end_date"""', 'type': '"""datetime in iso format"""', 'description': '"""All results will be before this datetime"""'}), "(name='__end_date', type='datetime in iso format', description=\n    'All results will be before this datetime')\n", (2698, 2809), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((5896, 5923), 'datetime.timedelta', 'timedelta', ([], {'weeks': '(4 * months)'}), '(weeks=4 * months)\n', (5905, 5923), False, 'from datetime import datetime, timedelta\n'), ((6752, 6777), 'streamlit.spinner', 'st.spinner', (['"""Thinking..."""'], {}), "('Thinking...')\n", (6762, 6777), True, 'import streamlit as st\n'), ((6806, 6882), 'streamlit.session_state.chat_engine.chat', 'st.session_state.chat_engine.chat', (['prompt'], {'function_call': '"""query_engine_tool"""'}), "(prompt, function_call='query_engine_tool')\n", (6839, 6882), True, 'import streamlit as st\n'), ((6899, 6926), 'streamlit.write', 'st.write', (['response.response'], {}), '(response.response)\n', (6907, 6926), True, 'import streamlit as st\n'), ((7021, 7062), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (['message'], {}), '(message)\n', (7053, 7062), True, 'import streamlit as st\n')]