ollama + pgvector

requirements.txt

@@ -53,7 +53,7 @@ unstructured
 # unstructured[all-docs]
 mixpanel
 langchain-mistralai==0.1.9
-langchain_community
+langchain_community==0.2.7
 huggingface_hub
 datasets
 langchain_text_splitters
@@ -70,4 +70,6 @@ pinecone-client
 pystemmer
 langchain_groq
 langchain-google-genai
-langchain-google-vertexai
+langchain-google-vertexai
+langchain-ollama
+langchain_postgres

setup.py

@@ -85,7 +85,7 @@
         'unstructured[all-docs]',
         'mixpanel',
         'langchain-mistralai==0.1.9',
-        'langchain_community',
+        'langchain_community==0.2.7',
         'huggingface_hub',
         'datasets',
         'langchain_text_splitters',
@@ -101,7 +101,9 @@
         'setuptools',
         'langchain_groq',
         'langchain-google-genai',
-        'langchain-google-vertexai'
+        'langchain-google-vertexai',
+        'langchain-ollama',
+        'langchain_postgres',
         # other dependencies
     ],
 )

src/ragbuilder/executor.py

@@ -69,6 +69,9 @@
 from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI
 from langchain_google_genai import ChatGoogleGenerativeAI,GoogleGenerativeAIEmbeddings
 from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings
+from langchain_postgres.vectorstores import PGVector
+from langchain_ollama.llms import OllamaLLM
+from langchain_ollama import OllamaEmbeddings
 
 # import local modules
 from ragbuilder.langchain_module.retriever.retriever import *

src/ragbuilder/langchain_module/embedding_model/embedding.py

@@ -15,7 +15,7 @@ def getEmbedding(**kwargs):
             raise KeyError("The key 'embedding_model' is missing from the arguments.")
 
         embedding_model = kwargs['embedding_model']
-
+        embedding_model='Ollama:llama3'
         model_owner= embedding_model.split(":")[0]
         model= embedding_model.split(":")[1]
         # Validate the embedding model type
@@ -48,6 +48,10 @@ def getEmbedding(**kwargs):
             logger.info(f"Azure Embedding Invoked: {embedding_model}")
             code_string= f"""embedding = AzureOpenAIEmbeddings(model='{model}')"""
             import_string = f"""from langchain_openai import AzureOpenAIEmbeddings"""
+        elif model_owner == "Ollama":
+            logger.info(f"Azure Embedding Invoked: {embedding_model}")
+            code_string= f"""embedding = OllamaEmbeddings(model='{embedding_model}')"""
+            import_string = f"""from langchain_ollama import OllamaEmbeddings"""
         else:
             raise ValueError(f"Invalid LLM: {embedding_model}")
         return {'code_string':code_string,'import_string':import_string}

src/ragbuilder/langchain_module/llms/llmConfig.py

@@ -9,6 +9,7 @@
 def getLLM(**kwargs):
     logger.info("LLM Invoked")
     retrieval_model=kwargs['retrieval_model']
+    retrieval_model='Ollama:llama3'
     model_owner= retrieval_model.split(":")[0]
     model= retrieval_model.split(":")[1]
     if model_owner == "Groq":
@@ -39,6 +40,11 @@ def getLLM(**kwargs):
         logger.info(f"LLM Code Gen Invoked: {retrieval_model}")
         import_string = f"""from langchain_huggingface import HuggingFaceEndpoint"""
         code_string = f"""llm=HuggingFaceEndpoint(repo_id='{model}',huggingfacehub_api_token=os.environ.get('HUGGINGFACEHUB_API_TOKEN'))"""
+    elif model_owner == "Olama":
+        logger.info(f"LLM Code Gen Invoked: {retrieval_model}")
+        import_string = f"""from langchain_ollama.llms import OllamaLLM"""
+        code_string = f"""llm = OllamaLLM(model='{model}')"""
+
     else:
         raise ValueError(f"Invalid LLM: {kwargs['retrieval_model']}")
     return {'code_string':code_string,'import_string':import_string}

src/ragbuilder/langchain_module/vectordb/vectordb.py

@@ -2,6 +2,7 @@
 import logging
 from langchain_chroma import Chroma
 from langchain_community.vectorstores import FAISS, SingleStoreDB
+from langchain_postgres.vectorstores import PGVector
 from langchain_pinecone import PineconeVectorStore
 from pinecone import Pinecone, ServerlessSpec
 import uuid
@@ -76,6 +77,12 @@ def getVectorDB(db_type,embedding_model):
     time.sleep(1)
 c=PineconeVectorStore.from_documents(splits, embedding, index_name='{index_name}')"""
         import_string = f"""from langchain_pinecone import PineconeVectorStore"""
+    elif db_type == "pgvector":
+        code_string= f"""
+        connection = "os.environ.get('PGVECTOR_CONNECTION_STRING')"
+        collection_name = '{index_name}'
+        c = PGVector(embeddings=embedding,collection_name=collection_name,connection=connection,use_jsonb=True)"""
+        import_string = f"""from langchain_postgres.vectorstores import PGVector"""
     else:
         raise ValueError(f"Unsupported db_type: {db_type}. Supported types are singleStoreDB, 'chromaDB', 'pineconeDB' and 'faissDB'.")
     return {'code_string':code_string,'import_string':import_string}