Iterate

indexnetwork · Jun 2, 2024 · 5f08e2b · 5f08e2b
1 parent 82c51e9
commit 5f08e2b
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 66 deletions.
diff --git a/indexer/src/app/modules/agent.module.ts b/indexer/src/app/modules/agent.module.ts
@@ -222,12 +222,8 @@ export class Agent {
             },
           },
         ]),
-        question: (input) => {
-          if (input.question) {
-            Logger.log(input.question, 'ChatService:answerChain:inputQuestion');
-            return input.question;
-          }
-        },
+        question: (input) => input.question,
+        information: (input) => input.information,
         chat_history: (input) => input.chat_history,
       },
       {
@@ -240,25 +236,8 @@ export class Agent {
               const serialized = docs_with_metadata.join('\n');
               return serialized;
             },
-            question: (input) => {
-              if (input.question) {
-                Logger.log(
-                  input.question,
-                  'ChatService:answerChain:inputQuestion',
-                );
-                return input.question;
-              }
-            },
-            information: (input) => {
-              if (input.information) {
-                Logger.log(
-                  input.information,
-                  'ChatService:answerChain:inputInformation',
-                );
-                return input.inormation;
-              }
-            },
-            relation: (input) => '',
+            question: (input) => input.question,
+            information: (input) => input.information,
             chat_history: (input) => formatChatHistory(input.chat_history),
           },
           prompt as any,

diff --git a/indexer/src/app/modules/chroma.module.ts b/indexer/src/app/modules/chroma.module.ts
@@ -5,48 +5,53 @@ import { OpenAIEmbeddings } from '@langchain/openai';
 @Global()
 @Module({})
 export class ChromaModule {
-	static register (apiKey: string): DynamicModule {
-		return {
-			module: ChromaModule,
-			global: true,
-			providers: [
-				{
-					provide: 'CHROMA_DB',
-					useFactory: async (): Promise<Chroma> => {
-						try {
+  static register(apiKey: string): DynamicModule {
+    return {
+      module: ChromaModule,
+      global: true,
+      providers: [
+        {
+          provide: 'CHROMA_DB',
+          useFactory: async (): Promise<Chroma> => {
+            try {
+              if (!process.env.OPENAI_API_KEY)
+                throw new Error('OpenAI API key is required');
 
-                            if (!process.env.OPENAI_API_KEY) throw new Error('OpenAI API key is required');
-
-                            Logger.debug(process.env.CHROMA_URL, 'ChromaModule');
-                            Logger.debug(process.env.CHROMA_COLLECTION_NAME, 'ChromaModule');
+              Logger.debug(process.env.CHROMA_URL, 'ChromaModule');
+              Logger.debug(process.env.CHROMA_COLLECTION_NAME, 'ChromaModule');
 
-                            const vectorStore = await Chroma.fromExistingCollection(
-                                new OpenAIEmbeddings({
-                                    openAIApiKey: apiKey,
-                                    modelName: process.env.MODEL_EMBEDDING,
-                                }),
-                                { 
-                                    collectionName: process.env.CHROMA_COLLECTION_NAME,
-                                    url: process.env.CHROMA_URL,
-                                }
-                            );
-
-                            const ensure = await vectorStore.ensureCollection();
-                            Logger.debug(JSON.stringify(ensure), 'ChromaModule:ensureCollection');
+              const vectorStore = await Chroma.fromExistingCollection(
+                new OpenAIEmbeddings({
+                  openAIApiKey: apiKey,
+                  modelName: process.env.MODEL_EMBEDDING,
+                }),
+                {
+                  collectionName: process.env.CHROMA_COLLECTION_NAME,
+                  url: process.env.CHROMA_URL,
+                },
+              );
 
-                            const docCount = await vectorStore.collection.count();
-                            Logger.debug(`Loaded ${docCount} documents from ChromaDB`, 'ChromaModule');
+              const ensure = await vectorStore.ensureCollection();
+              Logger.debug(
+                JSON.stringify(ensure),
+                'ChromaModule:ensureCollection',
+              );
 
-                            return vectorStore;
-
-						} catch (e) {
-							Logger.error('ChromaDB connection failure:', e, 'ChromaDBModule');
-							throw new Error('CHROMADB_CONNECTION_FAILURE');
-						}
-					},
-				},
-			],
-			exports: ['CHROMA_DB'],
-		};
-	}
-}
+              const docCount = await vectorStore.collection.count();
+              Logger.debug(
+                `Loaded ${docCount} documents from ChromaDB`,
+                'ChromaModule',
+              );
+
+              return vectorStore;
+            } catch (e) {
+              Logger.error('ChromaDB connection failure:', e, 'ChromaDBModule');
+              throw new Error('CHROMADB_CONNECTION_FAILURE');
+            }
+          },
+        },
+      ],
+      exports: ['CHROMA_DB'],
+    };
+  }
+}
diff --git a/indexer/src/chat/schema/chat.schema.ts b/indexer/src/chat/schema/chat.schema.ts
@@ -73,6 +73,7 @@ export class RetrievalQuestionInput {
     description: 'Chain input',
     example: {
       question: 'Summarize my documents',
+      information: 'Here is a json',
       chat_history: [],
     },
   })