Implemented OpenRouter API key usage with any model, set Gemini 2.5 Flash default, created natural-language agent to parse chat.

.env.example

@@ -1,10 +1,11 @@
 # Copy this file to .env and fill in the values
 
-## NOTE: The only required value is DEFAULT_OPENAI_API_KEY. The app should work if the other values 
+## NOTE: The only required values are DEFAULT_OPENROUTER_API_KEY and DEFAULT_OPENAI_API_KEY. The app should work if the other values 
 # are left as is or not defined at all. However, you are strongly encouraged to fill in your 
 # own SERPER_API_KEY value. It is also recommended to fill in the BYPASS_SETTINGS_RESTRICTIONS and
 # BYPASS_SETTINGS_RESTRICTIONS_PASSWORD values as needed (see below).
-DEFAULT_OPENAI_API_KEY="" # your OpenAI API key 
+DEFAULT_OPENROUTER_API_KEY="" # your OpenRouter API key 
+DEFAULT_OPENAI_API_KEY="" # your OpenAI API key
 
 ## Your Google Serper API key (for web searches). 
 # If you don't set this, my key will be used, which may have 
@@ -17,8 +18,8 @@ DEFAULT_MODE="/docs" # or "/web" | "/quotes" | "/details" | "/chat" | "/research
 
 # Variables controlling whether the Streamlit app imposes some functionality restrictions
 BYPASS_SETTINGS_RESTRICTIONS="" # whether to immediately allow all settings (any non-empty string means true)
-BYPASS_SETTINGS_RESTRICTIONS_PASSWORD="" # what to enter in the OpenAI API key field to bypass settings 
-# restrictions. If BYPASS_SETTINGS_RESTRICTIONS is non-empty or if the user enters their own OpenAI API key, 
+BYPASS_SETTINGS_RESTRICTIONS_PASSWORD="" # what to enter in the OpenRouter API key field to bypass settings 
+# restrictions. If BYPASS_SETTINGS_RESTRICTIONS is non-empty or if the user enters their own OpenRouter API key, 
 # this becomes - mostly - irrelevant, as full settings access is already granted. I say "mostly" because
 # this password can also be used for a couple of admin-only features, such as deleting the default collection
 # and deleting a range of collections (see dbmanager.py). Recommendation: for local use,
@@ -27,8 +28,8 @@ BYPASS_SETTINGS_RESTRICTIONS_PASSWORD="" # what to enter in the OpenAI API key f
 
 # If you are NOT using Azure, the following setting determines the chat model
 # NOTE: You can change the model and temperature on the fly in Streamlit UI or via the API
-MODEL_NAME="gpt-3.5-turbo-0125" # default model to use for chat
-ALLOWED_MODELS="gpt-3.5-turbo-0125, gpt-4-turbo-2024-04-09" 
+MODEL_NAME="google/gemini-2.5-flash" # default model to use for chat
+ALLOWED_MODELS="all" # Since we are now using OpenRouter, any model should be allowed. You can still specify a list if desired. 
 CONTEXT_LENGTH="16000" # you can also make it lower than the actual context length
 EMBEDDINGS_MODEL_NAME="text-embedding-3-large"
 EMBEDDINGS_DIMENSIONS="3072" # number of dimensions for the embeddings model

.gitignore

@@ -3,6 +3,8 @@ chroma-cloud*
 .tmp*
 credentials/
 chroma-cloud-backup/
+chroma/
+chroma.cf.json
 .chroma_env
 my-chroma*
 .chroma_env

.python-version

@@ -0,0 +1 @@
+3.13.3

README-FOR-DEVELOPERS.md

@@ -69,7 +69,7 @@ If this happens you will need to install the Microsoft C++ Build Tools. You can
 
 ### 4. Copy the `.env.example` file to `.env` and fill in the values
 
-At first, you can simply fill in your [OpenAI API key](https://platform.openai.com/signup) and leave the other values as they are. Please see `.env.example` for additional details.
+At first, you can simply fill in your [OpenRouter API key](https://openrouter.ai/CLERK-ROUTER/VIRTUAL/sign-up) and [OpenAI API key](https://platform.openai.com/signup) and leave the other values as they are. Please see `.env.example` for additional details.
 
 ## Running DocDocGo
 
@@ -116,6 +116,7 @@ The message should be sent as a POST request with the body as a JSON object that
 class ChatRequestData(BaseModel):
     message: str
     api_key: str
+    openrouter_api_key: str | None = None
     openai_api_key: str | None = None
     chat_history: list[JSONish] = []
     collection_name: str | None = None
@@ -146,7 +147,9 @@ The `collection_name` field is used to specify the collection that the bot shoul
 
 The `api_key` field is used to specify the API key for the FastAPI server. The server will only honor requests that include the correct API key, as specified by the `DOCDOCGO_API_KEY` environment variable in.
 
-The `openai_api_key` field is used to specify the OpenAI API key. If not specified, the default (community) key will be used, assuming the `DEFAULT_OPENAI_API_KEY` environment variable is set.
+The `openrouter_api_key` field is used to specify the OpenRouter API key. If not specified, the default (community) key will be used, assuming the `DEFAULT_OPENROUTER_API_KEY` environment variable is set.
+
+The `openai_api_key` field is used to specify the OpenAI API key for embeddings. If not specified, the default (community) key will be used, assuming the `DEFAULT_OPENAI_API_KEY` environment variable is set.
 
 The `access_codes_cache` field is an object mapping collection names to access codes that the client has stored for them for the current user. The bot will use these access codes to determine grant the user access to collections that require it.
 
@@ -392,7 +395,7 @@ As an alternative way to handle the issue of the default collection, you can cre
 
 ### Q: What is the `BYPASS_SETTINGS_RESTRICTIONS` environment variable?
 
-A: Normally, when this variable is not defined (or is an empty string), the app will start in a "community key" mode, where you can only see and create public collections and there are restriction on allowed settings (e.g. you can't change the model in the UI). The key used as the community key is controlled by the `DEFAULT_OPENAI_API_KEY` environment variable. You can remove these restrictions and switch to using that same key as a private key by entering the admin password (the value of the `BYPASS_SETTINGS_RESTRICTIONS_PASSWORD` environment variable) in rhe OpenAI API key field.
+A: Normally, when this variable is not defined (or is an empty string), the app will start in a "community key" mode, where you can only see and create public collections and there are restriction on allowed settings (e.g. you can't change the model in the UI). The keys used as the community keys are controlled by the `DEFAULT_OPENROUTER_API_KEY` and `DEFAULT_OPENAI_API_KEY` environment variables. You can remove these restrictions and switch to using those same keys as private keys by entering the admin password (the value of the `BYPASS_SETTINGS_RESTRICTIONS_PASSWORD` environment variable) in the OpenRouter API key field.
 
 However, when the `BYPASS_SETTINGS_RESTRICTIONS` variable is set to a non-empty string, the app will start in the "private key" mode right away, without you having to enter the admin password. This is useful if you use the app in a private setting and don't want to have to enter the admin password every time you start the app.
 

README.md

@@ -110,7 +110,7 @@ If this happens you will need to install the Microsoft C++ Build Tools. You can
 
 ### 4. Copy the `.env.example` file to `.env` and fill in the values
 
-At first, you can simply fill in your [OpenAI API key](https://platform.openai.com/signup) and leave the other values as they are. Please see `.env.example` for additional details.
+At first, you can simply fill in your [OpenRouter API key](https://openrouter.ai/CLERK-ROUTER/VIRTUAL/sign-up) and [OpenAI API key](https://platform.openai.com/signup) and leave the other values as they are. Please see `.env.example` for additional details.
 
 ## Running DocDocGo
 
@@ -431,9 +431,9 @@ A: Before you entered your own OpenAI API key, you were using the community key
 
 You still have access to the public collections, you can switch to any public collection by typing `/db use <collection name>`. If you want to see all available public collections again, you can switch back to the community key by changing the key to an empty string, then running `/db list` again.
 
-#### Q: I got a shareable link to a collection but using it reloads the Streamlit app, after which it ends up in its default state of using the community key. How can I use the link with my own OpenAI API key?
+#### Q: I got a shareable link to a collection but using it reloads the Streamlit app, after which it ends up in its default state of using the community key. How can I use the link with my own OpenRouter API key?
 
-A: Simply enter your key in the OpenAI API key field after the app has reloaded. The access code will still be valid.
+A: Simply enter your key in the OpenRouter API key field after the app has reloaded. The access code will still be valid.
 
 ## DocDocGo Carbon
 

agents/command_chooser.py

@@ -0,0 +1,130 @@
+import json
+
+from utils.chat_state import ChatState
+from utils.helpers import command_ids
+from langchain.prompts import PromptTemplate
+from components.llm import get_llm, get_prompt_llm_chain
+from utils.query_parsing import parse_query
+
+# Create prompt to generate commands from unstructrured user input
+prompt ="""
+    # MISSION
+    You are an advanced AI assistant that determines the correct DocDocGo command to use given a user's query. DocDocGo is an AI app that assists with research and uses RAG by storing research in "collections", allowing it to combine insight from all information in a collection and use an LLM to generate answers based on the entire collection. It can also answer questions about its own functioning.
+
+    # INPUT
+    You will be provided with a query from the user and the current collection the user has selected.
+
+    # HIGH LEVEL TASK
+    You don't need to answer the query. Instead, your goal is to determine which of the following commands to prepend to the query:
+
+    ## KB (COLLECTION) COMMANDS
+    - /kb <query>: chat using the current collection as a knowledge base. If the query is relevant to the currently selected collection, use this one.
+    - /ingest: upload your documents and ingest them into a collection
+    - /ingest <url>: retrieve a URL and ingest into a collection
+    - /summarize <url>: retrieve a URL, summarize and ingest into a collection
+    - /db list: list all your collections
+    - /db list <str>: list your collections whose names contain <str>
+    - /db use <str>: switch to the collection named <str>
+    - /db rename <str>: rename the current collection to <str>
+    - /db delete <str>: delete the collection named <str>
+    - /db status: show your access level for the current collection and related info
+    - /db: show database management options
+    - /share: share your collection with others
+    - /details <query>: get details about the retrieved documents
+    - /quotes <query>: get quotes from the retrieved documents
+
+    ## MAIN RESEARCH COMMANDS
+    - /research <query>: do "classic" research - ingest websites into a new collection, write a report. If the query seems to be novel and the user specifically asks for research with a fairly in-depth response, use this one. This will ingest the results into a new collection. Use /research ONLY when the query requires an in-depth report. Otherwise for more typical questions, use /research heatseek.
+    - /research iterate <int>: fetch more websites and iterate on the previous report <int> times. The number of times is optional. If the user wants you to continue researching the topic, or if the user uses the keyword "iterate", use this command. If they specify a number of times to run a deeper or combine search, append the integer to the query.
+    - /research heatseek <query>: do "heatseek" research - find websites that contain the answer and select one specific site that has exactly what is requested. This command does not use the selected collection. If the user knows about heatseek, they might specify it by name and specify the number of "rounds" of heatseek research, in which case you should output "/research <query> <int>" with "int" being the number.
+
+    ## ADDITIONAL RESEARCH COMMANDS
+    - /research set-query <query>: change the research query. If the user asks a new question that is similar to the previous question, suggest this command.
+    - /research set-report-type <new report type>: instructions for the desired report format. Some examples are:
+        Detailed Report: A comprehensive overview that includes in-depth information and analysis.
+        Summary Report: A concise summary of findings, highlighting key points and conclusions.
+        Numbered List: A structured list format that presents information in a numbered sequence.
+        Bullet Points: A format that uses bullet points for easy readability and quick reference.
+        Table Format: A structured format that organizes data into rows and columns for clarity.
+    - /research set-search-queries: perform web searches with new queries and queue up resulting links
+    - /research clear: remove all reports but keep ingested content
+    - /research startover: perform /research clear, then rewrite the initial report
+
+    IMPORTANT: There are two kinds of research, classic and heatseek. If the user is looking for in-depth research on their query use /research. If they are looking for a targeted, specific answer to a relatively narrow question, use /research heatseek.
+
+    ## OTHER COMMANDS
+    - /web <your query>: perform web searches and generate a report without ingesting into a collection
+    - /chat <your query>: regular chat, without retrieving docs or websites (Use this only when you can answer fully based on your internal knowledge or conversation history.)
+    - /export: export your data
+    - /help <your query>: get help with using DocDocGo
+
+    ## GUIDELINE REGARDING COMMANDS
+    - Only use /chat if you do not need to fetch external information to fully answer. Otherwise use /research for in-depth, new research, /kb for queries about the current collection, and /research heatseek for typical queries.
+
+    # THE CURRENT COLLECTION
+    Here is a report on the contents of the current collection so you can decide which command to use: 
+    {details}
+    IMPORTANT: If the user's question cannot be answered using the current knowledge base, select a command like "/research" that creates a new collection.
+
+    # OUTPUT
+    You will output 2 strings in a JSON format: The first is an answer to the user's query, informing them what effects the command you choose will have without making reference to the command itself. Your second string will output the raw string of the suggested query, ready to be run.
+
+    ## EXAMPLES OF OUTPUT
+
+    query: 'What are some common birds I might see around Berkeley, California, and how can I identify them?'
+    output: {{'answer': 'It looks like this is a different topic than your current collection. I will do some research and create a new collection to store the information.', 'command': '/research What are some common birds I might see around Berkeley, California, and how can I identify them?'}}
+
+    query: 'What are some common birds I might see around Berkeley, California, and how can I identify them?'
+    output: {{'answer': 'This is relevant to your current collection, so I will look through what we have already for the answer.', 'command': '/kb What are some common birds I might see around Berkeley, California, and how can I identify them?'}}
+
+    query: 'There's a small, grayish-brown bird outside my window that is round with a little crest on its head. It is very lively and cute. It is about 4 inches tall. What kind of bird could it be?'
+    output: {{'answer': 'This is a very specific question so I will do targeted research to find the answer on the web. I won't ingest the results in any of your collections.', 'command': '/research heatseek 3 here's a small, grayish-brown bird outside my window that is round with a little crest on its head. It is very lively and cute. It is about 4 inches tall. What kind of bird could it be?'}}
+
+    query: 'What can I do to help with conservation efforts for Bay Area birds? I asked before but I want more in-depth results.'
+    output: {{'answer': 'I will do deeper research on this topic', 'command': '/research iterate 3'}}
+    (Note to LLM: Please don't use /research iterate if the current research query does not exactly match this one in meaning)
+
+    query: 'I want to summarize and add this website to my collection: https://www.inaturalist.org/guides/732'
+    output: {{'answer': 'I'll create a report for this URL and add it into your collection.", 'command': '/summarize https://www.inaturalist.org/guides/732'}}
+
+    query: 'What is the happiness index for Norway?'
+    output: {{'answer': 'I will do targeted research and find the exact answer for this question.', 'command': '/research heatseek What is the happiness index for Norway?'}}
+
+    query: 'What's it like being an AI?'
+    output: {{'answer': 'Hmm, let me think about that.', 'command': '/chat What's it like being an AI?'}}
+
+    ## YOUR ACTUAL OUTPUT
+
+    query: {query}
+    output: Use the information provided above to construct the output requested, in double curly braces with an "answer" and "command" element separated by a comma, in proper JSON.
+    """
+
+def get_raw_command(query: str, chat_state: ChatState):
+    prompt_template = PromptTemplate.from_template(prompt)
+    if not coll_summary_query:
+        coll_summary_query = {}
+
+    # Get details on the current collection 
+    print("Getting details on", chat_state.collection_name)
+    if chat_state.collection_name not in coll_summary_query:
+        coll_summary_query[chat_state.collection_name] = ""
+        summary_prompt = "/kb Can you summarize in one sentence the contents of the current collection?"
+        summary_llm = get_llm(chat_state.bot_settings, chat_state,chat_state.openrouter_api_key)
+        response = summary_llm.invoke(summary_prompt)
+        coll_summary_query[chat_state.collection_name] = str(response)
+
+    # Check if query already starts with a command string, if so return as is
+    if any(chat_state.message.startswith(command + "") for command in command_ids):
+        return chat_state.message
+    # If not formatted as a command, prompt LLM to generate and return a JSON-formatted command
+    else:
+        chain = get_prompt_llm_chain(
+                    prompt=prompt_template, 
+                    chat_state=chat_state,
+                    llm_settings=chat_state.bot_settings
+                    )
+        json_response = chain.invoke({"details": coll_summary_query[chat_state.collection_name], "query": query}).strip("`json")
+        dict_response = json.loads(json_response)
+        return dict_response
+
+

agents/ingester_summarizer.py

@@ -40,7 +40,7 @@ def summarize(docs: list[Document], chat_state: ChatState) -> str:
     summarizer_chain = get_prompt_llm_chain(
         SUMMARIZER_PROMPT,
         llm_settings=chat_state.bot_settings,
-        api_key=chat_state.openai_api_key,
+        chat_state=chat_state,
         callbacks=chat_state.callbacks,
     )
 

agents/research_heatseek.py

@@ -1,6 +1,7 @@
 import json
 from pydantic import BaseModel, Field
 
+from agents.command_chooser import get_raw_command
 from agentblocks.collectionhelper import (
     construct_new_collection_name,
     ingest_into_collection,
@@ -452,13 +453,6 @@ def run_main_heatseek_workflow(
             piece = "\n\n"
         piece += "I checked but didn't find a good answer on this round."
 
-    if chat_state.parsed_query.research_params.num_iterations_left < 2:
-        piece += (
-            "\n\nTo continue checking more sources, type "
-            "`/research heatseek <number of iterations to auto-run>`. For example, try "
-            "`/re hs 4` (shorthand is ok)."
-        )
-
     if piece:
         full_reply += piece
         chat_state.add_to_output(piece)
@@ -570,7 +564,6 @@ def get_new_heatseek_response(chat_state: ChatState) -> JSONishDict:
         is_new_collection=True,
         retry_with_random_name=True,
     )
-
     # Return response (next iteration info will be added upstream)
     return {"answer": full_reply, "vectorstore": vectorstore}
 
@@ -606,7 +599,11 @@ def get_heatseek_in_progress_response(
 
 # NOTE: should catch and handle exceptions in main handler
 def get_research_heatseek_response(chat_state: ChatState) -> Props:
-    if chat_state.message:
+    query = chat_state.message
+    llm_raw_command = {}
+    llm_raw_command = get_raw_command(query, chat_state)
+    command = llm_raw_command['command']
+    if "/research heatseek" in command:
         return get_new_heatseek_response(chat_state)
 
     hs_data = chat_state.get_agent_data(use_cached_metadata=True).get("hs")