Provide WatsonX/BAM models as langchain.chat_models

examples/user/langchain_rag_code.py

@@ -0,0 +1,145 @@
+import os, logging
+from dotenv import load_dotenv
+
+from genai.extensions.langchain import WatsonX
+
+## source https://python.langchain.com/docs/use_cases/question_answering/code_understanding
+
+## Loading
+from git import Repo
+from langchain.text_splitter import Language
+from langchain.document_loaders.generic import GenericLoader
+from langchain.document_loaders.parsers import LanguageParser
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+# Clone
+repo_path = "/root/jupyter/test/test_repo/langchain"
+data_root = "/root/jupyter/test/data/rag"
+
+dataset = 'LangChain'
+data_dir = os.path.join(data_root, dataset)
+max_docs = -1
+
+if os.path.exists(os.path.join(repo_path, 'docs')):
+    pass
+else:
+    logger.info(f'clone in to {repo_path}')
+    # os.rmdir(repo_path)
+    repo = Repo.clone_from("https://github.com/langchain-ai/langchain", to_path=repo_path)
+
+# Load
+loader = GenericLoader.from_filesystem(
+    repo_path+"/libs/test/langchain",
+    glob="**/*",
+    suffixes=[".py"],
+    parser=LanguageParser(language=Language.PYTHON, parser_threshold=500)
+)
+documents = loader.load()
+
+print(f'len(documents) = {len(documents)}')
+logger.info(f'len(documents) = {len(documents)}')
+
+## Splitting
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+python_splitter = RecursiveCharacterTextSplitter.from_language(language=Language.PYTHON, 
+                                                               chunk_size=2000, 
+                                                               chunk_overlap=200)
+texts = python_splitter.split_documents(documents)
+len(texts)
+
+## RetrievalQA
+
+
+from langchain.vectorstores import Chroma
+
+# from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.embeddings import SentenceTransformerEmbeddings
+
+if os.path.exists(data_dir):
+    logger.info(f'load chroma vectors from {data_dir}')
+    db = Chroma(
+        embedding_function=SentenceTransformerEmbeddings(),
+        persist_directory=data_dir
+    )
+else:
+    logger.info(f'It takes long time to populate Chroma db. Please be patient ...')
+    db = Chroma.from_documents(
+        texts, 
+        SentenceTransformerEmbeddings(),
+        persist_directory=data_dir)
+
+
+retriever = db.as_retriever(
+    search_type="mmr", # Also test "similarity"
+    search_kwargs={"k": 8},
+)
+
+## Chat
+
+from langchain.memory import ConversationSummaryMemory
+from langchain.chains import ConversationalRetrievalChain
+
+load_dotenv()
+api_key = os.getenv("GENAI_KEY", None)
+api_endpoint = os.getenv("GENAI_API", None)
+
+model_kwargs = {
+    "decoding_method": "greedy",
+    "max_new_tokens": 200,
+    "min_new_tokens": 1,
+    "stream": False,
+    "temperature": 0.7,
+    "repetition_penalty": 1.0,
+    "top_k": 50,
+    "top_p": 1,
+}
+
+llm = WatsonX(
+    model_name="ibm/falcon-40b-8lang-instruct", 
+    model_kwargs=model_kwargs,
+    watsonx_api_key=api_key,
+    watsonx_api_base=api_endpoint,
+    verbose=True
+    )
+
+memory = ConversationSummaryMemory(llm=llm,memory_key="chat_history",return_messages=True)
+qa = ConversationalRetrievalChain.from_llm(llm, retriever=retriever, memory=memory)
+
+question = "How can I initialize a ReAct agent?"
+print(f'question = {question}')
+logger.info(f'question = {question}')
+result = qa(question)
+logger.info(f'result["answer"] = {result["answer"]}')
+
+questions = [
+    "What is the class hierarchy?",
+    "What classes are derived from the Chain class?",
+    "What one improvement do you propose in code in relation to the class hierarchy for the Chain class?",
+]
+
+for question in questions:
+    result = qa(question)
+    print(f"-> **Question**: {question} \n")
+    print(f"**Answer**: {result['answer']} \n")
+
+
+output = '''
+-> **Question**: What is the class hierarchy?
+
+**Answer**:  (superclass, subclass, etc.)
+
+-> **Question**: What classes are derived from the Chain class?
+
+**Answer**:
+Answer not in context
+
+-> **Question**: What one improvement do you propose in code in relation to the class hierarchy for the Chain class?
+
+**Answer**:
+
+The Chain class is a generic class that can be used to represent a chain of other classes. It is not specific to any particular application domain.
+'''

examples/user/langchain_rag_qa.py

@@ -0,0 +1,98 @@
+import os, logging
+from dotenv import load_dotenv
+
+# from genai.extensions.langchain import LangChainInterface
+
+try:
+    from genai.extensions.langchain import WatsonX
+except ImportError:
+    raise ImportError('Could not import WatsonX')
+
+
+## source https://python.langchain.com/docs/use_cases/question_answering
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+# Load documents
+
+from langchain.document_loaders import WebBaseLoader
+loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
+
+# Split documents
+logger.info(f'web loader is loading content ...')
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+text_splitter = RecursiveCharacterTextSplitter(chunk_size = 500, chunk_overlap = 0)
+splits = text_splitter.split_documents(loader.load())
+
+# Embed and store splits
+logger.info(f'vector soter is populating ...')
+from langchain.vectorstores import Chroma
+#from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings import SentenceTransformerEmbeddings
+vectorstore = Chroma.from_documents(documents=splits,embedding=SentenceTransformerEmbeddings())
+retriever = vectorstore.as_retriever()
+
+question = "What are the approaches to Task Decomposition?"
+docs = vectorstore.similarity_search(question)
+logger.info(f'len(docs) = {len(docs)}')
+
+# Prompt 
+# https://smith.langchain.com/hub/rlm/rag-prompt
+
+from langchain import hub
+
+## for https://smith.langchain.com
+langChain_api_key = os.getenv("LangChain_API_KEY", None)
+
+rag_prompt = hub.pull("rlm/rag-prompt", api_key=langChain_api_key)
+
+logger.info(f'rag_prompt = {rag_prompt}\n')
+# LLM
+
+
+# llm = Model("ibm/falcon-40b-8lang-instruct", params=params, credentials=creds)
+
+load_dotenv()
+api_key = os.getenv("GENAI_KEY", None)
+api_endpoint = os.getenv("GENAI_API", None)
+
+model_kwargs = {
+    "decoding_method": "greedy",
+    "max_new_tokens": 200,
+    "min_new_tokens": 1,
+    "stream": False,
+    "temperature": 0.7,
+    "repetition_penalty": 1.0,
+    "top_k": 50,
+    "top_p": 1,
+}
+
+llm = WatsonX(
+    model_name="ibm/falcon-40b-8lang-instruct", 
+    model_kwargs=model_kwargs,
+    watsonx_api_key=api_key,
+    watsonx_api_base=api_endpoint,
+    verbose=True
+    )
+
+from langchain.retrievers.multi_query import MultiQueryRetriever
+
+logging.basicConfig()
+logging.getLogger().setLevel(logging.INFO)
+
+retriever_from_llm = MultiQueryRetriever.from_llm(retriever=vectorstore.as_retriever(),
+                                                  llm=llm)
+unique_docs = retriever_from_llm.get_relevant_documents(query=question)
+len(unique_docs)
+
+# RAG chain 
+
+from langchain.schema.runnable import RunnablePassthrough
+rag_chain = (
+    {"context": retriever, "question": RunnablePassthrough()} 
+    | rag_prompt 
+    | llm 
+)
+
+msg = rag_chain.invoke("What is Task Decomposition?")
+print(f'invork msg = {msg}')

src/genai/extensions/langchain/__init__.py

@@ -1,4 +1,11 @@
 from genai.extensions.langchain.llm import LangChainInterface
 from genai.extensions.langchain.prompt import PromptExtension
+from genai.extensions.langchain.watsonx import WatsonX
+from genai.extensions.langchain.bam import BAM
 
-__all__ = ["LangChainInterface", "PromptExtension"]
+__all__ = [
+    "LangChainInterface", 
+    "PromptExtension",
+    "WatsonX",
+    "BAM",
+    ]

src/genai/extensions/langchain/bam.py

@@ -0,0 +1,62 @@
+from genai.credentials import Credentials
+# from genai.model import GenAiException
+
+import os, logging
+
+from langchain.pydantic_v1 import Field, BaseModel
+
+from genai import Credentials
+from genai import Model
+
+from typing import (
+    Any,
+)
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["BAM"]
+
+
+class BAM(BaseModel):
+    model_name: str = Field(default="ibm/falcon-40b-8lang-instruct", alias="model")
+    """Model name to use."""
+    watsonx_api_key: str = None
+    watsonx_api_base: str = None
+    watsonx_organization: str = "ibm"
+    model_kwargs: dict[str, Any] = None
+    model: Any = None 
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
+
+    def create(self, messages: str = None, **kwargs: Any ) -> Any:
+
+        logging.debug(f'\tIn BAM.create\nmessages = {messages}\n')
+        prompts = []
+        for m in messages:
+            logging.debug(f'm.role = {m["role"]}')
+            prompts.append(m["content"])
+
+        logging.debug(f'\tIn BAM.create\messages = {messages}\n')
+        logging.debug(f'\tIn BAM.create\nprompts = {prompts}\n')
+        logging.debug(f'In - BAM.create - kwargs = {kwargs}')
+        creds = Credentials(self.watsonx_api_key, self.watsonx_api_base)
+
+        model = Model(self.model_name, params=self.model_kwargs, credentials=creds)
+
+        response = {}
+        res = {}
+        res["message"] = {}
+        response ["choices"] = []
+        result = model.generate(prompts)
+        for index, rsp in enumerate(result):
+
+            res["message"]["content"] = rsp.generated_text
+            res["message"]["role"] = messages[index]["role"]
+            response ["choices"].append(res)
+
+        return response
+