langchain/templates/cohere-librarian/cohere_librarian/rag.py
Bagatur 480626dc99
docs, community[patch], experimental[patch], langchain[patch], cli[pa… (#15412)
…tch]: import models from community

ran
```bash
git grep -l 'from langchain\.chat_models' | xargs -L 1 sed -i '' "s/from\ langchain\.chat_models/from\ langchain_community.chat_models/g"
git grep -l 'from langchain\.llms' | xargs -L 1 sed -i '' "s/from\ langchain\.llms/from\ langchain_community.llms/g"
git grep -l 'from langchain\.embeddings' | xargs -L 1 sed -i '' "s/from\ langchain\.embeddings/from\ langchain_community.embeddings/g"
git checkout master libs/langchain/tests/unit_tests/llms
git checkout master libs/langchain/tests/unit_tests/chat_models
git checkout master libs/langchain/tests/unit_tests/embeddings/test_imports.py
make format
cd libs/langchain; make format
cd ../experimental; make format
cd ../core; make format
```
2024-01-02 15:32:16 -05:00

17 lines
441 B
Python

from langchain.retrievers import CohereRagRetriever
from langchain_community.chat_models import ChatCohere
rag = CohereRagRetriever(llm=ChatCohere())
def get_docs_message(message):
docs = rag.get_relevant_documents(message)
message_doc = next(
(x for x in docs if x.metadata.get("type") == "model_response"), None
)
return message_doc.page_content
def librarian_rag(x):
return get_docs_message(x["message"])