Build Chatbots That Know Your Business with MongoDB and Cohere

1
!pip install --quiet datasets tqdm cohere pymongo

1
import os
2
import cohere
3
4
os.environ["COHERE_API_KEY"] = ""
5
co = cohere.Client(os.environ.get("COHERE_API_KEY"))
6
7
os.environ["HF_TOKEN"] = ""

1
import pandas as pd
2
from datasets import load_dataset
3
4
# Make sure you have an Hugging Face token(HF_TOKEN) in your development environemnt before running the code below
5
# How to get a token: https://huggingface.co/docs/hub/en/security-tokens
6
# https://huggingface.co/datasets/MongoDB/fake_tech_companies_market_reports
7
dataset = load_dataset("MongoDB/fake_tech_companies_market_reports", split="train", streaming=True)
8
dataset_df = dataset.take(100)
9
10
# Convert the dataset to a pandas dataframe
11
dataset_df = pd.DataFrame(dataset_df)
12
dataset_df.head(5)

1
# Data Preparation
2
def combine_attributes(row):
3
  combined = f"{row['company']} {row['sector']} "
4
5
  # Add reports information
6
  for report in row['reports']:
7
    combined += f"{report['year']} {report['title']} {report['author']} {report['content']} "
8
9
  # Add recent news information
10
  for news in row['recent_news']:
11
    combined += f"{news['headline']} {news['summary']} "
12
13
  return combined.strip()

1
# Add the new column 'combined_attributes'
2
dataset_df['combined_attributes'] = dataset_df.apply(combine_attributes, axis=1)

1
# Display the first few rows of the updated dataframe
2
dataset_df[['company', 'ticker', 'combined_attributes']].head()

1
from tqdm import tqdm
2
3
def get_embedding(text: str, input_type: str="search_document") -> list[float]:
4
  if not text.strip():
5
    print("Attempted to get embedding for empty text.")
6
    return []
7
8
  model = "embed-english-v3.0"
9
  response = co.embed(
10
    texts=[text],
11
    model=model,
12
    input_type=input_type, # Used for embeddings of search queries run against a vector DB to find relevant documents
13
    embedding_types=['float']
14
  )
15
16
  return response.embeddings.float[0]
17
18
# Apply the embedding function with a progress bar
19
tqdm.pandas(desc="Generating embeddings")
20
dataset_df["embedding"] = dataset_df['combined_attributes'].progress_apply(get_embedding)
21
22
print(f"We just computed {len(dataset_df['embedding'])} embeddings.")

1
dataset_df.head()

{
  "fields": [
    {
      "numDimensions": 1024,
      "path": "embedding",
      "similarity": "cosine",
      "type": "vector"
    }
  ]
}

1
import os
2
os.environ["MONGO_URI"] = ""

1
import pymongo
2
3
def get_mongo_client(mongo_uri):
4
  """Establish and validate connection to the MongoDB."""
5
6
  client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.rag.cohere_mongodb.python")
7
8
  # Validate the connection
9
  ping_result = client.admin.command('ping')
10
  if ping_result.get('ok') == 1.0:
11
    # Connection successful
12
    print("Connection to MongoDB successful")
13
    return client
14
  else:
15
    print("Connection to MongoDB failed")
16
  return None
17
18
MONGO_URI = os.environ["MONGO_URI"]
19
20
if not MONGO_URI:
21
  print("MONGO_URI not set in environment variables")
22
23
mongo_client = get_mongo_client(MONGO_URI)
24
25
DB_NAME = "asset_management_use_case"
26
COLLECTION_NAME = "market_reports"
27
28
db = mongo_client.get_database(DB_NAME)
29
collection = db.get_collection(COLLECTION_NAME)

1
# Delete any existing records in the collection
2
collection.delete_many({})

DeleteResult({'n': 63, 'electionId': ObjectId('7fffffff000000000000002b'), 'opTime': {'ts': Timestamp(1721913981, 63), 't': 43}, 'ok': 1.0, '$clusterTime': {'clusterTime': Timestamp(1721913981, 63), 'signature': {'hash': b'cU;+\xe3\xbdRc\t\x80\xad\x03\x16\x11\x18\xe6s\xebF\x01', 'keyId': 7353740577831124994}}, 'operationTime': Timestamp(1721913981, 63)}, acknowledged=True)

1
documents = dataset_df.to_dict('records')
2
collection.insert_many(documents)
3
4
print("Data ingestion into MongoDB completed")

1
def vector_search(user_query, collection):
2
  """
3
  Perform a vector search in the MongoDB collection based on the user query.
4
5
  Args:
6
  user_query (str): The user's query string.
7
  collection (MongoCollection): The MongoDB collection to search.
8
9
  Returns:
10
  list: A list of matching documents.
11
  """
12
13
  # Generate embedding for the user query
14
  query_embedding = get_embedding(user_query, input_type="search_query")
15
16
  if query_embedding is None:
17
    return "Invalid query or embedding generation failed."
18
19
  # Define the vector search pipeline
20
  vector_search_stage = {
21
    "$vectorSearch": {
22
      "index": "vector_index",
23
      "queryVector": query_embedding,
24
      "path": "embedding",
25
      "numCandidates": 150,  # Number of candidate matches to consider
26
      "limit": 5  # Return top 4 matches
27
    }
28
  }
29
30
  unset_stage = {
31
    "$unset": "embedding"  # Exclude the 'embedding' field from the results
32
  }
33
34
  project_stage = {
35
    "$project": {
36
      "_id": 0,  # Exclude the _id field
37
      "company": 1,  # Include the plot field
38
      "reports": 1,  # Include the title field
39
      "combined_attributes": 1, # Include the genres field
40
      "score": {
41
        "$meta": "vectorSearchScore"  # Include the search score
42
      }
43
    }
44
  }
45
46
  pipeline = [vector_search_stage, unset_stage, project_stage]
47
48
  # Execute the search
49
  results = collection.aggregate(pipeline)
50
  return list(results)

1
def rerank_documents(query: str, documents, top_n: int = 3):
2
    # Perform reranking with Cohere ReRank Model
3
    try:
4
        response = co.rerank(
5
            model="rerank-english-v3.0",
6
            query=query,
7
            documents=documents,
8
            top_n=top_n,
9
            rank_fields=["company", "reports", "combined_attributes"]
10
        )
11
12
        # Extract the top reranked documents
13
        top_documents_after_rerank = []
14
        for result in response.results:
15
            original_doc = documents[result.index]
16
            top_documents_after_rerank.append({
17
                'company': original_doc['company'],
18
                'combined_attributes': original_doc['combined_attributes'],
19
                'reports': original_doc['reports'],
20
                'vector_search_score': original_doc['score'],
21
                'relevance_score': result.relevance_score
22
            })
23
24
        return top_documents_after_rerank
25
26
    except Exception as e:
27
        print(f"An error occurred during reranking: {e}")
28
        return documents[:top_n]  # Return top N documents without reranking

1
import pprint
2
3
query = "What companies have negative market reports or negative sentiment that might deter from investment in the long term"
4
5
get_knowledge = vector_search(query, collection)
6
pd.DataFrame(get_knowledge).head()

1
reranked_documents = rerank_documents(query, get_knowledge)
2
pd.DataFrame(reranked_documents).head()

1
def format_documents_for_chat(documents):
2
  return [
3
    {
4
        "company": doc['company'],
5
        # "reports": doc['reports'],
6
        "combined_attributes": doc['combined_attributes']
7
    }
8
    for doc in documents
9
  ]
10
11
# Generating response with Cohere Command R
12
response = co.chat(
13
  message=query,
14
  documents=format_documents_for_chat(reranked_documents),
15
  model="command-r-plus",
16
  temperature=0.3
17
)
18
19
print("Final answer:")
20
print(response.text)

1
for cite in response.citations:
2
  print(cite)

start=122 end=145 text='GreenEnergy Corp (GRNE)' document_ids=['doc_0']
start=151 end=161 text='Challenges' document_ids=['doc_0']
start=173 end=231 text='solid financial performance and a positive market position' document_ids=['doc_0']
start=266 end=322 text='volatile political environment and rising trade tensions' document_ids=['doc_0']
start=337 end=384 text='increased tariffs and supply chain disruptions.' document_ids=['doc_0']
start=390 end=409 text='Regulatory Scrutiny' document_ids=['doc_0']
start=428 end=474 text='under scrutiny for its data handling practices' document_ids=['doc_0']
start=484 end=547 text='concerns about potential privacy breaches and ethical dilemmas.' document_ids=['doc_0']
start=552 end=578 text='BioEngineering Corp (BENC)' document_ids=['doc_1']
start=584 end=602 text='Regulatory Hurdles' document_ids=['doc_1']
start=617 end=667 text='delays in obtaining approvals for certain products' document_ids=['doc_1']
start=675 end=707 text='stringent healthcare regulations' document_ids=['doc_1']
start=725 end=740 text='time-to-market.' document_ids=['doc_1']
start=745 end=780 text='Reimbursement and Pricing Pressures' document_ids=['doc_1']
start=787 end=808 text='healthcare costs rise' document_ids=['doc_1']
start=827 end=864 text='carefully navigate pricing strategies' document_ids=['doc_1']
start=868 end=908 text='balance accessibility and profitability.' document_ids=['doc_1']
start=913 end=946 text='Research and Development Expenses' document_ids=['doc_1']
start=973 end=1009 text='significant increase in R&D expenses' document_ids=['doc_1']
start=1043 end=1083 text='maintain a competitive pricing strategy.' document_ids=['doc_1']
start=1088 end=1113 text='QuantumSensor Corp (QSCP)' document_ids=['doc_2']
start=1119 end=1143 text='Supply Chain Disruptions' document_ids=['doc_2']
start=1162 end=1181 text='supply chain issues' document_ids=['doc_2']
start=1189 end=1240 text='global logistics problems and geopolitical tensions' document_ids=['doc_2']
start=1252 end=1276 text='production and delivery.' document_ids=['doc_2']
start=1281 end=1300 text='Regulatory Scrutiny' document_ids=['doc_2']
start=1319 end=1380 text='under scrutiny for its data collection and handling practices' document_ids=['doc_2']
start=1387 end=1426 text='potential privacy and ethical concerns.' document_ids=['doc_2']
start=1431 end=1461 text='Technical Workforce Challenges' document_ids=['doc_2']
start=1465 end=1528 text='Attracting and retaining skilled talent in a competitive market' document_ids=['doc_2']

1
from typing import Dict, Optional, List
2
3
class CohereChat:
4
5
    def __init__(self, cohere_client, system: str = "", database: str = "cohere_chat",
6
                 main_collection: str = "main_collection", history_params: Optional[Dict[str, str]] = None):
7
      self.co = cohere_client
8
      self.system = system
9
      self.history_params = history_params or {}
10
11
      # Use the connection string from history_params
12
      self.client = pymongo.MongoClient(self.history_params.get('connection_string', 'mongodb://localhost:27017/'))
13
14
      # Use the database parameter
15
      self.db = self.client[database]
16
17
      # Use the main_collection parameter
18
      self.main_collection = self.db[main_collection]
19
20
      # Use the history_collection from history_params, or default to "chat_history"
21
      self.history_collection = self.db[self.history_params.get('history_collection', 'chat_history')]
22
23
      # Use the session_id from history_params, or default to "default_session"
24
      self.session_id = self.history_params.get('session_id', 'default_session')
25
26
    def add_to_history(self, message: str, prefix: str = ""):
27
      self.history_collection.insert_one({
28
        'session_id': self.session_id,
29
        'message': message,
30
        'prefix': prefix
31
      })
32
33
    def get_chat_history(self) -> List[Dict[str, str]]:
34
      history = self.history_collection.find({'session_id': self.session_id}).sort('_id', 1)
35
      return [{"role": "user" if item['prefix'] == "USER" else "chatbot", "message": item['message']} for item in history]
36
37
    def rerank_documents(self, query: str, documents: List[Dict], top_n: int = 3) -> List[Dict]:
38
      rerank_docs = [
39
          {
40
            'company': doc['company'],
41
            'combined_attributes': doc['combined_attributes']
42
          }
43
          for doc in documents
44
          if doc['combined_attributes'].strip()
45
      ]
46
47
      if not rerank_docs:
48
          print("No valid documents to rerank.")
49
          return []
50
51
      try:
52
          response = self.co.rerank(
53
              query=query,
54
              documents=rerank_docs,
55
              top_n=top_n,
56
              model="rerank-english-v3.0",
57
              rank_fields=["company", "combined_attributes"]
58
          )
59
60
          top_documents_after_rerank = [
61
              {
62
                  'company': rerank_docs[result.index]['company'],
63
                  'combined_attributes': rerank_docs[result.index]['combined_attributes'],
64
                  'relevance_score': result.relevance_score
65
              }
66
              for result in response.results
67
          ]
68
69
          print(f"\nHere are the top {top_n} documents after rerank:")
70
          for doc in top_documents_after_rerank:
71
              print(f"== {doc['company']} (Relevance: {doc['relevance_score']:.4f})")
72
73
          return top_documents_after_rerank
74
75
      except Exception as e:
76
          print(f"An error occurred during reranking: {e}")
77
          return documents[:top_n]
78
79
    def format_documents_for_chat(self, documents: List[Dict]) -> List[Dict]:
80
      return [
81
        {
82
          "company": doc['company'],
83
          "combined_attributes": doc['combined_attributes']
84
        }
85
        for doc in documents
86
      ]
87
88
    def send_message(self, message: str, vector_search_func) -> str:
89
      self.add_to_history(message, "USER")
90
91
      # Perform vector search
92
      search_results = vector_search_func(message, self.main_collection)
93
94
      # Rerank the search results
95
      reranked_documents = self.rerank_documents(message, search_results)
96
97
      # Format documents for chat
98
      formatted_documents = self.format_documents_for_chat(reranked_documents)
99
100
      # Generate response using Cohere chat
101
      response = self.co.chat(
102
        chat_history=self.get_chat_history(),
103
        message=message,
104
        documents=formatted_documents,
105
        model="command-r-plus",
106
        temperature=0.3
107
      )
108
109
      result = response.text
110
      self.add_to_history(result, "CHATBOT")
111
112
      print("Final answer:")
113
      print(result)
114
115
      print("\nCitations:")
116
      for cite in response.citations:
117
        print(cite)
118
119
      return result
120
121
    def show_history(self):
122
      history = self.history_collection.find({'session_id': self.session_id}).sort('_id', 1)
123
      for item in history:
124
        print(f"{item['prefix']}: {item['message']}")
125
        print("-------------------------")

1
# Initialize CohereChat
2
chat = CohereChat(
3
    co,
4
    system="You are a helpful assistant taking on the role of an Asset Manager focused on tech companies.",
5
    database=DB_NAME,
6
    main_collection=COLLECTION_NAME,
7
    history_params={
8
        'connection_string': MONGO_URI,
9
        'history_collection': "chat_history",
10
        'session_id': 2
11
    }
12
)
13
14
# Send a message
15
response = chat.send_message("What is the best investment to make why?", vector_search)

Here are the top 3 documents after rerank:
== EcoTech Innovations (Relevance: 0.0001)
== GreenEnergy Systems (Relevance: 0.0001)
== QuantumComputing Inc (Relevance: 0.0000)
Final answer:
I am an AI assistant and cannot comment on what the single "best" investment is. However, I have found some companies that have been recommended as "Buy" investments in the documents provided. 
## EcoTech Innovations (ETIN)
EcoTech Innovations is a leading provider of sustainable technology solutions, specializing in renewable energy and environmentally friendly products. In 2023 and 2024, ETIN demonstrated solid financial performance, innovative capabilities, and a growing market presence, making it an attractive investment opportunity for those interested in the sustainable technology sector. 
## GreenEnergy Systems (GESY)
GreenEnergy Systems is a leading provider of renewable energy solutions, offering solar and wind power technologies, energy storage systems, and smart grid solutions. In 2023 and 2024, GESY reported strong financial performance, innovative product developments, and a solid market position, positioning it well for future growth in the renewable energy sector. 
## QuantumComputing Inc. (QCMP)
QuantumComputing Inc. is a leading developer of quantum computing software and solutions, aiming to revolutionize computing tasks across industries. In 2023 and 2024, QCMP demonstrated strong financial performance, innovative product offerings, and a growing market presence, making it an attractive investment opportunity in the rapidly growing quantum computing industry. 
Please note that these recommendations are based on specific reports and may not consider all factors. It is always advisable to conduct thorough research and consult professional advice before making any investment decisions.
Citations:
start=148 end=153 text='"Buy"' document_ids=['doc_0', 'doc_1', 'doc_2']
start=198 end=224 text='EcoTech Innovations (ETIN)' document_ids=['doc_0']
start=250 end=302 text='leading provider of sustainable technology solutions' document_ids=['doc_0']
start=320 end=375 text='renewable energy and environmentally friendly products.' document_ids=['doc_0']
start=379 end=383 text='2023' document_ids=['doc_0']
start=388 end=392 text='2024' document_ids=['doc_0']
start=412 end=439 text='solid financial performance' document_ids=['doc_0', 'doc_1']
start=441 end=464 text='innovative capabilities' document_ids=['doc_0']
start=472 end=495 text='growing market presence' document_ids=['doc_0', 'doc_1']
start=572 end=602 text='sustainable technology sector.' document_ids=['doc_0']
start=608 end=634 text='GreenEnergy Systems (GESY)' document_ids=['doc_1']
start=660 end=706 text='leading provider of renewable energy solutions' document_ids=['doc_1']
start=717 end=801 text='solar and wind power technologies, energy storage systems, and smart grid solutions.' document_ids=['doc_1']
start=805 end=809 text='2023' document_ids=['doc_1']
start=814 end=818 text='2024' document_ids=['doc_1']
start=834 end=862 text='strong financial performance' document_ids=['doc_1']
start=864 end=895 text='innovative product developments' document_ids=['doc_1']
start=903 end=924 text='solid market position' document_ids=['doc_1']
start=971 end=995 text='renewable energy sector.' document_ids=['doc_1']
start=1001 end=1029 text='QuantumComputing Inc. (QCMP)' document_ids=['doc_2']
start=1057 end=1118 text='leading developer of quantum computing software and solutions' document_ids=['doc_2']
start=1130 end=1178 text='revolutionize computing tasks across industries.' document_ids=['doc_2']
start=1182 end=1186 text='2023' document_ids=['doc_2']
start=1191 end=1195 text='2024' document_ids=['doc_2']
start=1215 end=1243 text='strong financial performance' document_ids=['doc_2']
start=1245 end=1273 text='innovative product offerings' document_ids=['doc_2']
start=1281 end=1304 text='growing market presence' document_ids=['doc_2']
start=1360 end=1403 text='rapidly growing quantum computing industry.' document_ids=['doc_2']

1
# Show chat history
2
chat.show_history()

USER: What is the best investment to make why?
-------------------------
CHATBOT: I am an AI assistant and therefore cannot comment on what the single "best" investment is. However, I can tell you about some companies that have been recommended as "Buy" investments in the documents provided. 
## CloudInfra Systems (CISY)
CloudInfra Systems is a leading provider of cloud computing solutions, offering infrastructure-as-a-service (IaaS) and platform-as-a-service (PaaS) to businesses worldwide. In 2023, CISY demonstrated strong financial performance and product innovation, making it an attractive investment opportunity. 
## VirtualReality Systems (VRSY)
VirtualReality Systems is a leading provider of virtual reality hardware and software solutions. In 2023, VRSY reported strong financial performance, innovative product developments, and strategic partnerships, positioning it well in a rapidly growing and competitive market. 
## BioTech Innovations (BTCI)
BioTech Innovations is a leading biotechnology company specializing in healthcare solutions and innovative medicines. In 2023, BTCI demonstrated solid financial growth, product innovations, and a strengthened market position, making it an attractive investment option for long-term growth prospects. 
Please note that these recommendations are based on specific reports and may not consider all factors. It is always advisable to conduct thorough research and consult professional advice before making any investment decisions.
-------------------------
USER: What is the best investment to make why?
-------------------------
CHATBOT: I am an AI assistant and therefore cannot comment on what the single "best" investment is. However, I can provide you with some companies that have been recommended as "Buy" investments in the documents provided. 
## CloudInfra Systems (CISY)
CloudInfra Systems is a leading provider of cloud computing solutions, offering infrastructure-as-a-service (IaaS) and platform-as-a-service (PaaS) to businesses worldwide. In 2023, CISY demonstrated strong financial performance and product innovation, making it an attractive investment opportunity. 
## VirtualReality Systems (VRSY)
VirtualReality Systems is a leading provider of virtual reality hardware and software solutions. In 2023, VRSY reported strong financial performance, innovative product developments, and strategic partnerships, positioning it well in a rapidly growing and competitive market. 
## BioTech Innovations (BTCI)
BioTech Innovations is a leading biotechnology company specializing in healthcare solutions and innovative medicines. In 2023, BTCI demonstrated solid financial growth, product innovations, and a strengthened market position, making it an attractive investment option for long-term growth prospects. 
Please note that these recommendations are based on specific reports and may not consider all factors. It is always advisable to conduct thorough research and consult professional advice before making any investment decisions.
-------------------------
USER: What is the best investment to make why?
-------------------------
CHATBOT: I am an AI assistant and cannot comment on what the single "best" investment is. However, I can provide information on companies that have been recommended as "Buy" investments in the documents provided. 
## CloudInfra Systems (CISY)
CloudInfra Systems is a leading provider of cloud computing solutions, offering infrastructure-as-a-service (IaaS) and platform-as-a-service (PaaS) to a diverse range of businesses. In 2023, CISY demonstrated strong financial performance and product innovation, positioning it well in the competitive cloud market. 
## VirtualReality Systems (VRSY)
VirtualReality Systems is a leading provider of virtual reality hardware and software solutions. In 2023, VRSY reported robust financial results, innovative product developments, and strategic partnerships, making it a solid investment choice for those with a long-term investment horizon. 
## BioTech Innovations (BTCI)
BioTech Innovations is a leading biotechnology company specializing in healthcare solutions and innovative medicines. In 2023 and 2024, BTCI demonstrated solid financial growth, product innovations, and an improved market position, making it an attractive investment opportunity for long-term growth. 
Please note that these recommendations are based on specific reports and may not consider all factors. It is always advisable to conduct thorough research and consult professional advice before making any investment decisions.
-------------------------
USER: What is the best investment to make why?
-------------------------
CHATBOT: I am an AI assistant and cannot comment on what the single "best" investment is. However, I have found some companies that have been recommended as "Buy" investments in the documents provided. 
## EcoTech Innovations (ETIN)
EcoTech Innovations is a leading provider of sustainable technology solutions, specializing in renewable energy and environmentally friendly products. In 2023 and 2024, ETIN demonstrated solid financial performance, innovative capabilities, and a growing market presence, making it an attractive investment opportunity for those interested in the sustainable technology sector. 
## GreenEnergy Systems (GESY)
GreenEnergy Systems is a leading provider of renewable energy solutions, offering solar and wind power technologies, energy storage systems, and smart grid solutions. In 2023 and 2024, GESY reported strong financial performance, innovative product developments, and a solid market position, positioning it well for future growth in the renewable energy sector. 
## QuantumComputing Inc. (QCMP)
QuantumComputing Inc. is a leading developer of quantum computing software and solutions, aiming to revolutionize computing tasks across industries. In 2023 and 2024, QCMP demonstrated strong financial performance, innovative product offerings, and a growing market presence, making it an attractive investment opportunity in the rapidly growing quantum computing industry. 
Please note that these recommendations are based on specific reports and may not consider all factors. It is always advisable to conduct thorough research and consult professional advice before making any investment decisions.
-------------------------