code: import asyncio
  import json
  import os
  import tempfile
  import time
  from functools import lru_cache
  from logging import getLogger
  from pathlib import Path
  
  from fastapi import APIRouter, Request, status
  from fastapi.encoders import jsonable_encoder
  from fastapi.responses import HTMLResponse
  from typing import List, Dict, Any
  from pydantic import Field, validator
  
  # This is here to satisfy runtime import needs 
  # that pyinstaller appears to miss
  
  from llama_index.node_parser import SentenceSplitter
  from llama_index.schema import TextNode, NodeRelationship, RelatedNodeInfo, MetadataMode, NodeWithScore
  from llama_index.callbacks import CallbackManager, LlamaDebugHandler, OpenInferenceCallbackHandler
  from llama_index.embeddings import OpenAIEmbedding, OllamaEmbedding
  from llama_index.indices.query.query_transform import HyDEQueryTransform
  from llama_index.query_pipeline import QueryPipeline
  from llama_index.llms import OpenAI, Ollama
  from llama_index.llms.base import BaseLLM
  from llama_index.vector_stores.chroma import ChromaVectorStore
  from llama_index import LLMPredictor, PromptTemplate, VectorStoreIndex, Document, StorageContext, ServiceContext, download_loader
  from llama_index.callbacks import CallbackManager, LlamaDebugHandler
  from llama_index.embeddings import OpenAIEmbedding
  from llama_index.llms import OpenAI
  from llama_index.vector_stores.chroma import ChromaVectorStore
  from llama_index import VectorStoreIndex, Document, StorageContext, ServiceContext, download_loader
  from llama_index.callbacks import CallbackManager, LlamaDebugHandler
  from llama_index.indices.query.query_transform.base import DecomposeQueryTransform
  from llama_index import ServiceContext
  from llama_index.postprocessor import CohereRerank
  from llama_index.response_synthesizers import TreeSummarize
  from llama_index.postprocessor import PrevNextNodePostprocessor, LLMRerank
  from llama_index.storage.docstore import SimpleDocumentStore
  from llama_index.query_pipeline import CustomQueryComponent, InputKeys, OutputKeys
  from llama_index.postprocessor.types import BaseNodePostprocessor
  from llama_index.vector_stores.types import BasePydanticVectorStore
  
  from llama_index.indices.vector_store.retrievers import VectorIndexAutoRetriever, VectorIndexRetriever
  
  from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo
  
  from snowflake import SnowflakeGenerator
  
  from service.dependencies import (
      TANA_NODE,
      TANA_TEXT,
      LlamaindexAsk,
      TanaNodeMetadata,
  )
  from service.endpoints.chroma import get_collection, get_tana_nodes_by_id
  
  from service.endpoints.topics import TanaDocument, extract_topics, is_reference_content, tana_node_ids_from_text
  from service.llamaindex import DecomposeQueryWithNodeContext, WidenNodeWindowPostProcessor, create_index, get_index
  from service.tana_types import TanaDump
  
  logger = getLogger()
  snowflakes = SnowflakeGenerator(42)
  
  router = APIRouter()
  
  minutes = 1000 * 60
  
  # TODO: Add header support throughout so we can pass Tana API key and OpenAPI Key as headers
  # NOTE: we already have this in the main.py middleware wrapper, but it would be better
  # to do it here for OpenAPI spec purposes.
  # x_tana_api_token: Annotated[str | None, Header()] = None
  # x_openai_api_key: Annotated[str | None, Header()] = None
  
  
  # enrich our retriever with knowledge of our metadata
  def get_auto_retriever(index:VectorStoreIndex):
    vector_store_info = VectorStoreInfo(
      content_info="My Tana Notebook. Comprises many Tana nodes with text and metadata fields.",
      metadata_info=[
          MetadataInfo(
              name="category",
              type="str",
              description=(
                  "One of TANA_NODE or TANA_TEXT\n"
                  "TANA_NODE means that this is a top-level topic in my Tana notebook\n"
                  "TANA_TEXT means this is detailed information as part of a topic, identfied by topic_id metadata.\n"
                  "Do NOT use category to query the index. Only use category to enrich your understanding of the result.\n"
                  "DO NOT reference category in your responses.\n"
              ),
          ),
          MetadataInfo(
              name="topic_id",
              type="str",
              description=(
                  "Identifies the Tana Notebook Node that this text is part of. Should be used as a reference to the notebook entry.\n"
                  "Only use topic_id to query the index when you want a single specific node by reference.\n"
                  "You can use topic_id when referencing a Tana Notebook Node in your responses.\n"
              ),
          ),
          MetadataInfo(
              name="tana_id",
              type="str",
              description=(
                  "The Tana Notebook Node for this piece of text. Should be used a reference to the notebook entry.\n"
                  "Only use topic_id to query the index when you want a single specific node by reference.\n"
                  "You can use tana_id when referencing a Tana Notebook Node in your responses.\n"
              ),
          ),
          MetadataInfo(
              name="supertag",
              type="str",
              description=(
                  "One or more optional GENERAL semantic ontology tags for this Tana Notebook Node.\n"
                  "Delimited by spaces (NOT a LIST. Do not use IN operator to test membership)\n"
                  "Example: \n"
                  "{ supertag:  #task #topic #person #meeting }\n"
                  "Do NOT use supertags to query the index. Only use supertags to enrich your understanding of the result.\n"
              ),
          ),
        ],
    )
  
    # THIS doesn't work at all well with GPT 3
    # and only works sometimes with GPT4. Problem is that it becomes fixated on the 
    # use of metadata to filter results, overly constraining relevance.
    # retriever = VectorIndexAutoRetriever(
    #     index, 
    #     vector_store_info=vector_store_info, 
    #     similarity_top_k=10
    # )
  
    retriever = VectorIndexRetriever(index=index, similarity_top_k=10)
    return retriever
  
  
  @router.post("/llamaindex/ask", response_class=HTMLResponse, tags=["research"])
  def llamaindex_ask(req: LlamaindexAsk, model:str):
    '''Ask a question of the Llamaindex and return the top results
    '''
  
    (index, service_context, vector_store, llm) = get_index(model=model)
  
    query_engine=index.as_query_engine(similarity_top_k=20, stream=False)
  
    logger.info(f'Querying LLamaindex with {req.query}')
    response = query_engine.query(req.query)
    return str(response)
  
  
  summary_tmpl = PromptTemplate(
      "You are an expert Q&A system that is trusted around the world.\n"
      "TASK\n"
      "Summarize the following CONTEXT in order to best answer the QUERY.\n"
      "Answer the QUERY using the provided CONTEXT information, and not prior knowledge.\n"
      "Some rules to follow:\n"
      "1. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\n"
      "2. The CONTEXT contais references to many Tana Notebook Nodes. Nodes have both metadata and text content\n"
      "3. Whenever your summary needs to reference Tana Notebook Nodes from the CONTEXT, use proper Tana node reference format as follows:\n"
      "  the characters '[[' + '^' + tana_id metadata and then the characters ']]'.\n"
      "  E.g. to reference the Tana context node titled 'Recipe for making icecream' with tana_id: xghysd76 use this format:\n"
      "    [[^xghysd76]]\n"
      "5. Try to avoid making many redundant references to the same Tana node in your summary. Use footnote style if you really need to do this.\n"
      "\n"
      "QUERY: {query_str}\n"
      "-----\n"
      "CONTEXT:\n"
      "{context_str}\n"
      "END_CONTEXT\n"
      "-----\n"
    )
  
  #TODO: Move model out of POST body and into query params perhaps?
  @router.post("/llamaindex/research", response_class=HTMLResponse, tags=["research"])
  def llama_ask_custom_pipeline(req: LlamaindexAsk, model:str):
    '''Research a question using Llamaindex and return the top results.'''
    (index, service_context, storage_context, llm) = get_index(model, observe=True)
  
    logger.info(f'Researching LLamaindex with {req.query}')
  
    # first, build up a set of research questions
    decompose_transform = DecomposeQueryWithNodeContext(llm=llm)
    p1 = QueryPipeline(chain=[decompose_transform])
    questions = p1.run(query=req.query)
    
   
    retriever = get_auto_retriever(index)
    # and preprocess the result nodes to make use of next/previous
    prevnext = WidenNodeWindowPostProcessor(storage_context=storage_context, num_nodes=5, mode="both")
    summarizer = TreeSummarize(summary_template=summary_tmpl, service_context=service_context)
  
    # for each question, do a fetch against Chroma to find potentially relevant nodes
    results = []
    for question in questions:
      if question == '':
        continue
  
      logger.info(f'Question: {question}')
      # use our metadata aware auto-retriever to fetch from Chroma
      q1 = QueryPipeline(chain=[retriever, prevnext])
      nodes = q1.run(input=question)
      # nodes = retriever.retrieve(question)
      # logger.info(f'Nodes:\n{nodes}')
  
      # clean up the redudant metadata (TANA_TEXT node metadata is less useful here)
      new_nodes = []
      if nodes:
        for node in nodes:
          new_node = node
          if node.metadata['category'] == TANA_TEXT:
            # copy the outer NodeWithScore and the inner TextNode objects
            new_text_node = TextNode(**node.node.dict())
            # wipe out the metadata
            new_text_node.metadata = {}
            new_node = NodeWithScore(node=new_text_node, score=node.score)
  
          new_nodes.append(new_node)
  
      research = '\n'.join([node.get_content(metadata_mode=MetadataMode.LLM) for node in new_nodes])
      logger.info(f'Nodes:\n{research}')
  
      # tailor the summarizer prompt
      sum_result = summarizer.as_query_component().run_component(nodes=new_nodes, query_str=question)
      summary = sum_result['output'].response
      logger.info(f'Summary:\n{summary}')
  
      result = {'question': question,
                'answers': nodes,
                'summary': summary}
      results.append(result)
  
    # now build up the context from the result nodes
    context = []
    for result in results:
      question = result['question']
      answer = result['answers']
      summary = result['summary']
      context.append(f'QUESTION: {question}\n')
  
      #context.append('RESEARCH:\n')
      # TODO: instead of dumping all nodes into the primary context
      # we should prepare an answer to each question and then use that
      # node:TextNode
      # for node in answer:
      #   context.append(node.get_content(metadata_mode=MetadataMode.LLM)+'\n')
      
      context.append('ANSWER:\n')
      context.append(summary+'\n')
  
      context.append('\n')
  
    # now combine all that research 
    prompt_tmpl = PromptTemplate(
      "You are an expert Q&A system that is trusted around the world.\n"
      "Always answer the question using the provided context information, and not prior knowledge.\n"
      "Some rules to follow:\n"
      "1. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\n"
      "2. You will be given CONTEXT information in the form of one or more related QUESTIONS and the ANSWERS to those questions.\n"
      "3. For each ANSWER, there may be many Tana Notebook Nodes. Nodes have both metadata and text content\n"
      "4. Whenever your response needs to reference Tana Notebook Nodes from the context, use proper Tana node reference format as follows:\n"
      "  the characters '[[' + '^' + tana_id metadata and then the characters ']]'.\n"
      "  E.g. to reference the Tana context node titled 'Recipe for making icecream' with tana_id: xghysd76 use this format:\n"
      "    [[^xghysd76]]\n"
      "5. Try to avoid making many redundant references to the same Tana node in your response. Use footnote style if you really need to do this.\n"
      "\n"
      "QUERY: {query}\n"
      "-----\n"
      "CONTEXT:\n"
      "{context}\n"
      "END_CONTEXT\n"
      "-----\n"
    )
  
    p2 = QueryPipeline(chain=[prompt_tmpl, llm])
    response = p2.run(query=req.query, context='\n'.join(context))
    return response.message.content
  
  # attempt to paralleize non-async code
  # see https://github.com/tiangolo/fastapi/discussions/6347
  lock = asyncio.Lock()
  
  
apis: llama_index.response_synthesizers.TreeSummarize llama_index.vector_stores.types.MetadataInfo llama_index.query_pipeline.QueryPipeline llama_index.PromptTemplate llama_index.indices.vector_store.retrievers.VectorIndexRetriever llama_index.schema.NodeWithScore
extract_api: [((2834, 2845), 'logging.getLogger', 'getLogger', ([], {}), '()\n', (2843, 2845), False, 'from logging import getLogger\n'), ((2859, 2881), 'snowflake.SnowflakeGenerator', 'SnowflakeGenerator', (['(42)'], {}), '(42)\n', (2877, 2881), False, 'from snowflake import SnowflakeGenerator\n'), ((2892, 2903), 'fastapi.APIRouter', 'APIRouter', ([], {}), '()\n', (2901, 2903), False, 'from fastapi import APIRouter, Request, status\n'), ((6478, 7521), 'llama_index.PromptTemplate', 'PromptTemplate', (['"""You are an expert Q&A system that is trusted around the world.\nTASK\nSummarize the following CONTEXT in order to best answer the QUERY.\nAnswer the QUERY using the provided CONTEXT information, and not prior knowledge.\nSome rules to follow:\n1. Avoid statements like \'Based on the context, ...\' or \'The context information ...\' or anything along those lines.\n2. The CONTEXT contais references to many Tana Notebook Nodes. Nodes have both metadata and text content\n3. Whenever your summary needs to reference Tana Notebook Nodes from the CONTEXT, use proper Tana node reference format as follows:\n  the characters \'[[\' + \'^\' + tana_id metadata and then the characters \']]\'.\n  E.g. to reference the Tana context node titled \'Recipe for making icecream\' with tana_id: xghysd76 use this format:\n    [[^xghysd76]]\n5. Try to avoid making many redundant references to the same Tana node in your summary. Use footnote style if you really need to do this.\n\nQUERY: {query_str}\n-----\nCONTEXT:\n{context_str}\nEND_CONTEXT\n-----\n"""'], {}), '(\n    """You are an expert Q&A system that is trusted around the world.\nTASK\nSummarize the following CONTEXT in order to best answer the QUERY.\nAnswer the QUERY using the provided CONTEXT information, and not prior knowledge.\nSome rules to follow:\n1. Avoid statements like \'Based on the context, ...\' or \'The context information ...\' or anything along those lines.\n2. The CONTEXT contais references to many Tana Notebook Nodes. Nodes have both metadata and text content\n3. Whenever your summary needs to reference Tana Notebook Nodes from the CONTEXT, use proper Tana node reference format as follows:\n  the characters \'[[\' + \'^\' + tana_id metadata and then the characters \']]\'.\n  E.g. to reference the Tana context node titled \'Recipe for making icecream\' with tana_id: xghysd76 use this format:\n    [[^xghysd76]]\n5. Try to avoid making many redundant references to the same Tana node in your summary. Use footnote style if you really need to do this.\n\nQUERY: {query_str}\n-----\nCONTEXT:\n{context_str}\nEND_CONTEXT\n-----\n"""\n    )\n', (6492, 7521), False, 'from llama_index import LLMPredictor, PromptTemplate, VectorStoreIndex, Document, StorageContext, ServiceContext, download_loader\n'), ((12205, 12219), 'asyncio.Lock', 'asyncio.Lock', ([], {}), '()\n', (12217, 12219), False, 'import asyncio\n'), ((5916, 5970), 'llama_index.indices.vector_store.retrievers.VectorIndexRetriever', 'VectorIndexRetriever', ([], {'index': 'index', 'similarity_top_k': '(10)'}), '(index=index, similarity_top_k=10)\n', (5936, 5970), False, 'from llama_index.indices.vector_store.retrievers import VectorIndexAutoRetriever, VectorIndexRetriever\n'), ((6243, 6265), 'service.llamaindex.get_index', 'get_index', ([], {'model': 'model'}), '(model=model)\n', (6252, 6265), False, 'from service.llamaindex import DecomposeQueryWithNodeContext, WidenNodeWindowPostProcessor, create_index, get_index\n'), ((7999, 8029), 'service.llamaindex.get_index', 'get_index', (['model'], {'observe': '(True)'}), '(model, observe=True)\n', (8008, 8029), False, 'from service.llamaindex import DecomposeQueryWithNodeContext, WidenNodeWindowPostProcessor, create_index, get_index\n'), ((8162, 8200), 'service.llamaindex.DecomposeQueryWithNodeContext', 'DecomposeQueryWithNodeContext', ([], {'llm': 'llm'}), '(llm=llm)\n', (8191, 8200), False, 'from service.llamaindex import DecomposeQueryWithNodeContext, WidenNodeWindowPostProcessor, create_index, get_index\n'), ((8208, 8250), 'llama_index.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[decompose_transform]'}), '(chain=[decompose_transform])\n', (8221, 8250), False, 'from llama_index.query_pipeline import QueryPipeline\n'), ((8412, 8503), 'service.llamaindex.WidenNodeWindowPostProcessor', 'WidenNodeWindowPostProcessor', ([], {'storage_context': 'storage_context', 'num_nodes': '(5)', 'mode': '"""both"""'}), "(storage_context=storage_context, num_nodes=5,\n    mode='both')\n", (8440, 8503), False, 'from service.llamaindex import DecomposeQueryWithNodeContext, WidenNodeWindowPostProcessor, create_index, get_index\n'), ((8515, 8592), 'llama_index.response_synthesizers.TreeSummarize', 'TreeSummarize', ([], {'summary_template': 'summary_tmpl', 'service_context': 'service_context'}), '(summary_template=summary_tmpl, service_context=service_context)\n', (8528, 8592), False, 'from llama_index.response_synthesizers import TreeSummarize\n'), ((10727, 11820), 'llama_index.PromptTemplate', 'PromptTemplate', (['"""You are an expert Q&A system that is trusted around the world.\nAlways answer the question using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Avoid statements like \'Based on the context, ...\' or \'The context information ...\' or anything along those lines.\n2. You will be given CONTEXT information in the form of one or more related QUESTIONS and the ANSWERS to those questions.\n3. For each ANSWER, there may be many Tana Notebook Nodes. Nodes have both metadata and text content\n4. Whenever your response needs to reference Tana Notebook Nodes from the context, use proper Tana node reference format as follows:\n  the characters \'[[\' + \'^\' + tana_id metadata and then the characters \']]\'.\n  E.g. to reference the Tana context node titled \'Recipe for making icecream\' with tana_id: xghysd76 use this format:\n    [[^xghysd76]]\n5. Try to avoid making many redundant references to the same Tana node in your response. Use footnote style if you really need to do this.\n\nQUERY: {query}\n-----\nCONTEXT:\n{context}\nEND_CONTEXT\n-----\n"""'], {}), '(\n    """You are an expert Q&A system that is trusted around the world.\nAlways answer the question using the provided context information, and not prior knowledge.\nSome rules to follow:\n1. Avoid statements like \'Based on the context, ...\' or \'The context information ...\' or anything along those lines.\n2. You will be given CONTEXT information in the form of one or more related QUESTIONS and the ANSWERS to those questions.\n3. For each ANSWER, there may be many Tana Notebook Nodes. Nodes have both metadata and text content\n4. Whenever your response needs to reference Tana Notebook Nodes from the context, use proper Tana node reference format as follows:\n  the characters \'[[\' + \'^\' + tana_id metadata and then the characters \']]\'.\n  E.g. to reference the Tana context node titled \'Recipe for making icecream\' with tana_id: xghysd76 use this format:\n    [[^xghysd76]]\n5. Try to avoid making many redundant references to the same Tana node in your response. Use footnote style if you really need to do this.\n\nQUERY: {query}\n-----\nCONTEXT:\n{context}\nEND_CONTEXT\n-----\n"""\n    )\n', (10741, 11820), False, 'from llama_index import LLMPredictor, PromptTemplate, VectorStoreIndex, Document, StorageContext, ServiceContext, download_loader\n'), ((11960, 11999), 'llama_index.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[prompt_tmpl, llm]'}), '(chain=[prompt_tmpl, llm])\n', (11973, 11999), False, 'from llama_index.query_pipeline import QueryPipeline\n'), ((8876, 8918), 'llama_index.query_pipeline.QueryPipeline', 'QueryPipeline', ([], {'chain': '[retriever, prevnext]'}), '(chain=[retriever, prevnext])\n', (8889, 8918), False, 'from llama_index.query_pipeline import QueryPipeline\n'), ((3533, 3945), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""category"""', 'type': '"""str"""', 'description': '"""One of TANA_NODE or TANA_TEXT\nTANA_NODE means that this is a top-level topic in my Tana notebook\nTANA_TEXT means this is detailed information as part of a topic, identfied by topic_id metadata.\nDo NOT use category to query the index. Only use category to enrich your understanding of the result.\nDO NOT reference category in your responses.\n"""'}), '(name=\'category\', type=\'str\', description=\n    """One of TANA_NODE or TANA_TEXT\nTANA_NODE means that this is a top-level topic in my Tana notebook\nTANA_TEXT means this is detailed information as part of a topic, identfied by topic_id metadata.\nDo NOT use category to query the index. Only use category to enrich your understanding of the result.\nDO NOT reference category in your responses.\n"""\n    )\n', (3545, 3945), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((4101, 4452), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""topic_id"""', 'type': '"""str"""', 'description': '"""Identifies the Tana Notebook Node that this text is part of. Should be used as a reference to the notebook entry.\nOnly use topic_id to query the index when you want a single specific node by reference.\nYou can use topic_id when referencing a Tana Notebook Node in your responses.\n"""'}), '(name=\'topic_id\', type=\'str\', description=\n    """Identifies the Tana Notebook Node that this text is part of. Should be used as a reference to the notebook entry.\nOnly use topic_id to query the index when you want a single specific node by reference.\nYou can use topic_id when referencing a Tana Notebook Node in your responses.\n"""\n    )\n', (4113, 4452), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((4568, 4900), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""tana_id"""', 'type': '"""str"""', 'description': '"""The Tana Notebook Node for this piece of text. Should be used a reference to the notebook entry.\nOnly use topic_id to query the index when you want a single specific node by reference.\nYou can use tana_id when referencing a Tana Notebook Node in your responses.\n"""'}), '(name=\'tana_id\', type=\'str\', description=\n    """The Tana Notebook Node for this piece of text. Should be used a reference to the notebook entry.\nOnly use topic_id to query the index when you want a single specific node by reference.\nYou can use tana_id when referencing a Tana Notebook Node in your responses.\n"""\n    )\n', (4580, 4900), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((5016, 5403), 'llama_index.vector_stores.types.MetadataInfo', 'MetadataInfo', ([], {'name': '"""supertag"""', 'type': '"""str"""', 'description': '"""One or more optional GENERAL semantic ontology tags for this Tana Notebook Node.\nDelimited by spaces (NOT a LIST. Do not use IN operator to test membership)\nExample: \n{ supertag:  #task #topic #person #meeting }\nDo NOT use supertags to query the index. Only use supertags to enrich your understanding of the result.\n"""'}), '(name=\'supertag\', type=\'str\', description=\n    """One or more optional GENERAL semantic ontology tags for this Tana Notebook Node.\nDelimited by spaces (NOT a LIST. Do not use IN operator to test membership)\nExample: \n{ supertag:  #task #topic #person #meeting }\nDo NOT use supertags to query the index. Only use supertags to enrich your understanding of the result.\n"""\n    )\n', (5028, 5403), False, 'from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n'), ((9472, 9523), 'llama_index.schema.NodeWithScore', 'NodeWithScore', ([], {'node': 'new_text_node', 'score': 'node.score'}), '(node=new_text_node, score=node.score)\n', (9485, 9523), False, 'from llama_index.schema import TextNode, NodeRelationship, RelatedNodeInfo, MetadataMode, NodeWithScore\n')]