code: import asyncio
  import io
  import json
  import logging
  import os
  import tempfile
  import time
  from logging import getLogger
  
  from fastapi import APIRouter, Request, status
  from fastapi.encoders import jsonable_encoder
  from typing import List, Tuple
  
  # This is here to satisfy runtime import needs 
  # that pyinstaller appears to miss
  
  from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo
  from llama_index import Document, VectorStoreIndex
  from llama_index import Document
  
  from snowflake import SnowflakeGenerator
  
  from service.dependencies import (
      TANA_NODE,
      TANA_TEXT,
      ChromaRequest,
      LlamaRequest,
      TanaNodeMetadata,
      capture_logs,
  )
  
  from service.endpoints.chroma import chroma_upsert
  from service.endpoints.topics import TanaDocument, extract_topic_from_context, extract_topics
  from service.llamaindex import create_index, get_index
  from service.tana_types import TanaDump
  from service.txntimer import txntimer
  
  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
  
  # TODO: change this to remove LLamaindex and simply go directly to ChromaDB
  
  
  async def load_chromadb_from_topics(topics:List[TanaDocument], model:str, observe=False):
    '''Load the topic index from the topic array directly.'''
  
    logger.info('Building ChromaDB vectors from nodes')
  
    index_nodes = []
    # loop through all the topics and create a Document for each
    for topic in topics:
      (doc_node, text_nodes) = document_from_topic(topic)
      index_nodes.append(doc_node)
      index_nodes.extend(text_nodes)
  
    logger.info(f'Gathered {len(index_nodes)} tana nodes')
    logger.info("Preparing storage context")
  
    for node in index_nodes:
      logger.info(f'Node {node.node_id} {node.metadata}')
      chroma_req = ChromaRequest(context=node.text, nodeId=node.node_id, model=model)
      upsert = await chroma_upsert(chroma_req)
      
    logger.info("ChromaDB populated and ready")
    return index_nodes
  
  def load_index_from_topics(topics:List[TanaDocument], model:str, observe=False):
    '''Load the topic index from the topic array directly.'''
  
    logger.info('Building llama_index nodes')
  
    index_nodes = []
    # loop through all the topics and create a Document for each
    for topic in topics:
      (doc_node, text_nodes) = document_from_topic(topic)
      index_nodes.append(doc_node)
      index_nodes.extend(text_nodes)
  
    logger.info(f'Gathered {len(index_nodes)} tana nodes')
    logger.info("Preparing storage context")
  
    index = create_index(model, observe, index_nodes)
    logger.info("Llamaindex populated and ready")
    return index
  
  
  def document_from_topic(topic) -> Tuple[Document, List[TextNode]]:
    '''Load a single topic into the index_nodes list.'''
    text_nodes = []
  
    metadata = {
      'category': TANA_NODE,
      'supertag': ' '.join(['#' + tag for tag in topic.tags]),
      'title': topic.name,
      'tana_id': topic.id,
      'document_id': topic.id,
      }
    
    if topic.fields:
      # get all the fields as metdata as well
      fields = set([field.name for field in topic.fields])
      for field_name in fields:
        metadata[field_name] = ' '.join([field.value for field in topic.fields if field.name == field_name])
  
    # what other props do we need to create?
    # document = Document(id_=topic.id, text=topic.name) # type: ignore
    # we only add the ffirst line and fields to the document payload
    # anything else and we blow out the token limits (and cost a lot!)
    text = topic.content[0][2]
    document_node = Document(doc_id=topic.id, text=text) # first line only
    document_node.metadata = metadata
  
    # # make a note of the document in our nodes list
    # index_nodes.append(document_node)
  
    # now iterate all the remaining topic.content and create a node for each
    # each of these is simply a string, being the name of a tana child node
    # but with [[]name^id]] reference syntax used for references
    # TODO: make these tana_nodes richer structurally
    # TODO: use actual tana node id here perhaps?
    previous_text_node = None
    if len(topic.content) > 30:
      logger.warning(f'Large topic {topic.id} with {len(topic.content)} children')
  
    # process all the child content records...
    for (content_id, is_ref, tana_element) in topic.content[1:]:
      content_metadata = TanaNodeMetadata(
        category=TANA_TEXT,
        title=topic.name,
        topic_id=topic.id,
        # TODO: ? 'supertag': ' '.join(['#' + tag for tag in topic.tags]),
        # text gets added below...
      )
      
      # wire up the tana_node as an index_node with the text as the payload
      if is_ref:
        current_text_node = TextNode(text=tana_element)
        current_text_node.metadata['tana_ref_id'] = content_id
      else:
        current_text_node = TextNode(id_=content_id, text=tana_element)
        current_text_node.metadata['tana_id'] = content_id
  
      current_text_node.metadata = content_metadata.model_dump()
  
      # check if this is a reference node and add additional metadata
      # TODO: backport this to chroma upsert...?
  
      current_text_node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=document_node.node_id)
      # wire up next/previous
      if previous_text_node:
        current_text_node.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(node_id=previous_text_node.node_id)
        previous_text_node.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(node_id=current_text_node.node_id)
  
      text_nodes.append(current_text_node)
      previous_text_node = current_text_node
    
    return (document_node, text_nodes)
  
  # attempt to paralleize non-async code
  # see https://github.com/tiangolo/fastapi/discussions/6347
  lock = asyncio.Lock()
  
  # Note: accepts ?model= query param
  @router.post("/llamaindex/upsert", status_code=status.HTTP_204_NO_CONTENT, tags=["preload"])
  async def llama_upsert(request: Request, req:LlamaRequest, model:str="openai", observe:bool=False):
    '''Upserts a single Tana node context into the Llama index.
    '''
    async with lock:
      tana_id = req.nodeId
      tana_context = req.context
      (index, _, _, _)  = get_index(model, observe)
      result = extract_topic_from_context(tana_id, tana_context)
      (document, text_nodes) = document_from_topic(result)
      # remove then add back
      # index.delete_nodes([document.node_id] + [node.node_id for node in text_nodes])
      # strip '0' ids from node list
      id_list = [document.node_id] + [node.node_id for node in text_nodes if node.node_id != '0']
      for node_id in id_list:
        index.delete_ref_doc(node_id) 
      
      # filter out the '0' id nodes
      text_nodes = [node for node in text_nodes if node.node_id != '0']
      node_list = [document] + text_nodes
      index.insert_nodes(node_list)
  
      #TODO: consider returning some kind of completion confirmation payload with statidtics
      
      return None
  
  
  # Note: accepts ?model= query param
  @router.post("/llamaindex/preload", tags=["preload"])
  async def llama_preload(request: Request, tana_dump:TanaDump, model:str="openai"):
    '''Accepts a Tana dump JSON payload and builds the index from it.
    Uses the topic extraction code from the topics endpoint to build
    an object tree in memory, then loads that into ChromaDB via LLamaIndex.
    
    Returns a list of log messages from the process.
    '''
    async with lock:
      messages = []
      async with capture_logs(logger) as logs:
        result = await extract_topics(tana_dump, 'JSON') # type: ignore
        logger.info('Extracted topics from Tana dump')
  
        # save output to a temporary file
        with tempfile.TemporaryDirectory() as tmp:
          path = os.path.join(tmp, 'topics.json')
          logger.info(f'Saving topics to {path}')
          # use path
          with open(path, "w") as f:
            json_result = jsonable_encoder(result)
            f.write(json.dumps(json_result))
          
          logger.info('Loading index ...')
        # don't use file anymore...
        # load_index_from_file(path)
        # load directly from in-memory representation
        await load_chromadb_from_topics(result, model=model)
        # load_index_from_topics(result, model=model)
      
        # logger.info(f'Deleted temp file {path}')
        messages = logs.getvalue()
      return messages
  
  
apis: llama_index.schema.TextNode llama_index.schema.RelatedNodeInfo llama_index.Document
extract_api: [((975, 986), 'logging.getLogger', 'getLogger', ([], {}), '()\n', (984, 986), False, 'from logging import getLogger\n'), ((1001, 1023), 'snowflake.SnowflakeGenerator', 'SnowflakeGenerator', (['(42)'], {}), '(42)\n', (1019, 1023), False, 'from snowflake import SnowflakeGenerator\n'), ((1034, 1045), 'fastapi.APIRouter', 'APIRouter', ([], {}), '()\n', (1043, 1045), False, 'from fastapi import APIRouter, Request, status\n'), ((5959, 5973), 'asyncio.Lock', 'asyncio.Lock', ([], {}), '()\n', (5971, 5973), False, 'import asyncio\n'), ((2838, 2879), 'service.llamaindex.create_index', 'create_index', (['model', 'observe', 'index_nodes'], {}), '(model, observe, index_nodes)\n', (2850, 2879), False, 'from service.llamaindex import create_index, get_index\n'), ((3831, 3867), 'llama_index.Document', 'Document', ([], {'doc_id': 'topic.id', 'text': 'text'}), '(doc_id=topic.id, text=text)\n', (3839, 3867), False, 'from llama_index import Document\n'), ((2125, 2191), 'service.dependencies.ChromaRequest', 'ChromaRequest', ([], {'context': 'node.text', 'nodeId': 'node.node_id', 'model': 'model'}), '(context=node.text, nodeId=node.node_id, model=model)\n', (2138, 2191), False, 'from service.dependencies import TANA_NODE, TANA_TEXT, ChromaRequest, LlamaRequest, TanaNodeMetadata, capture_logs\n'), ((4597, 4670), 'service.dependencies.TanaNodeMetadata', 'TanaNodeMetadata', ([], {'category': 'TANA_TEXT', 'title': 'topic.name', 'topic_id': 'topic.id'}), '(category=TANA_TEXT, title=topic.name, topic_id=topic.id)\n', (4613, 4670), False, 'from service.dependencies import TANA_NODE, TANA_TEXT, ChromaRequest, LlamaRequest, TanaNodeMetadata, capture_logs\n'), ((5392, 5438), 'llama_index.schema.RelatedNodeInfo', 'RelatedNodeInfo', ([], {'node_id': 'document_node.node_id'}), '(node_id=document_node.node_id)\n', (5407, 5438), False, 'from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo\n'), ((6371, 6396), 'service.llamaindex.get_index', 'get_index', (['model', 'observe'], {}), '(model, observe)\n', (6380, 6396), False, 'from service.llamaindex import create_index, get_index\n'), ((6410, 6459), 'service.endpoints.topics.extract_topic_from_context', 'extract_topic_from_context', (['tana_id', 'tana_context'], {}), '(tana_id, tana_context)\n', (6436, 6459), False, 'from service.endpoints.topics import TanaDocument, extract_topic_from_context, extract_topics\n'), ((2211, 2236), 'service.endpoints.chroma.chroma_upsert', 'chroma_upsert', (['chroma_req'], {}), '(chroma_req)\n', (2224, 2236), False, 'from service.endpoints.chroma import chroma_upsert\n'), ((4922, 4949), 'llama_index.schema.TextNode', 'TextNode', ([], {'text': 'tana_element'}), '(text=tana_element)\n', (4930, 4949), False, 'from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo\n'), ((5047, 5090), 'llama_index.schema.TextNode', 'TextNode', ([], {'id_': 'content_id', 'text': 'tana_element'}), '(id_=content_id, text=tana_element)\n', (5055, 5090), False, 'from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo\n'), ((5561, 5612), 'llama_index.schema.RelatedNodeInfo', 'RelatedNodeInfo', ([], {'node_id': 'previous_text_node.node_id'}), '(node_id=previous_text_node.node_id)\n', (5576, 5612), False, 'from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo\n'), ((5677, 5727), 'llama_index.schema.RelatedNodeInfo', 'RelatedNodeInfo', ([], {'node_id': 'current_text_node.node_id'}), '(node_id=current_text_node.node_id)\n', (5692, 5727), False, 'from llama_index.schema import BaseNode, TextNode, NodeRelationship, RelatedNodeInfo\n'), ((7617, 7637), 'service.dependencies.capture_logs', 'capture_logs', (['logger'], {}), '(logger)\n', (7629, 7637), False, 'from service.dependencies import TANA_NODE, TANA_TEXT, ChromaRequest, LlamaRequest, TanaNodeMetadata, capture_logs\n'), ((7668, 7701), 'service.endpoints.topics.extract_topics', 'extract_topics', (['tana_dump', '"""JSON"""'], {}), "(tana_dump, 'JSON')\n", (7682, 7701), False, 'from service.endpoints.topics import TanaDocument, extract_topic_from_context, extract_topics\n'), ((7822, 7851), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (7849, 7851), False, 'import tempfile\n'), ((7875, 7907), 'os.path.join', 'os.path.join', (['tmp', '"""topics.json"""'], {}), "(tmp, 'topics.json')\n", (7887, 7907), False, 'import os\n'), ((8034, 8058), 'fastapi.encoders.jsonable_encoder', 'jsonable_encoder', (['result'], {}), '(result)\n', (8050, 8058), False, 'from fastapi.encoders import jsonable_encoder\n'), ((8077, 8100), 'json.dumps', 'json.dumps', (['json_result'], {}), '(json_result)\n', (8087, 8100), False, 'import json\n')]