format updates + search added to RAG instead

ai_generate.py

@@ -11,14 +11,16 @@ import google.generativeai as genai
 import anthropic
 from langchain_community.document_loaders import PyMuPDFLoader
 from langchain_community.document_loaders import TextLoader
+from langchain_core.documents import Document
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
 )
+from langchain.schema import StrOutputParser
 from langchain_community.vectorstores import Chroma
 from langchain_text_splitters import CharacterTextSplitter
 from langchain import hub
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
+from langchain_core.runnables import RunnablePassthrough, RunnableMap
 from langchain.chains import RetrievalQA
 from langchain_groq import ChatGroq
 from langchain_openai import ChatOpenAI
@@ -46,147 +48,119 @@ claude_client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
 temperature = 1.0
 max_tokens = 2048
 
-rag_llms = {
-    "LLaMA 3": ChatGroq(
-        temperature=temperature,
-        max_tokens=max_tokens,
-        model_name="llama3-70b-8192",
-    ),
-    "OpenAI GPT 4o Mini": ChatOpenAI(
-        temperature=temperature,
-        max_tokens=max_tokens,
-        model_name="gpt-4o-mini",
-    ),
-    "OpenAI GPT 4o": ChatOpenAI(
-        temperature=temperature,
-        max_tokens=max_tokens,
-        model_name="gpt-4o",
-    ),
-    "OpenAI GPT 4": ChatOpenAI(
-        temperature=temperature,
-        max_tokens=max_tokens,
-        model_name="gpt-4-turbo",
-    ),
-    "Gemini 1.5 Pro": ChatGoogleGenerativeAI(temperature=temperature, max_tokens=max_tokens, model="gemini-1.5-pro"),
-    "Claude Sonnet 3.5": ChatAnthropic(
-        temperature=temperature,
-        max_tokens=max_tokens,
-        model_name="claude-3-5-sonnet-20240620",
-    ),
+llm_model_translation = {
+    "LLaMA 3": "llama3-70b-8192",
+    "OpenAI GPT 4o Mini": "gpt-4o-mini",
+    "OpenAI GPT 4o": "gpt-4o",
+    "OpenAI GPT 4": "gpt-4-turbo",
+    "Gemini 1.5 Pro": "gemini-1.5-pro",
+    "Claude Sonnet 3.5": "claude-3-5-sonnet-20240620",
 }
 
+llm_classes = {
+    "llama3-70b-8192": ChatGroq,
+    "gpt-4o-mini": ChatOpenAI,
+    "gpt-4o": ChatOpenAI,
+    "gpt-4-turbo": ChatOpenAI,
+    "gemini-1.5-pro": ChatGoogleGenerativeAI,
+    "claude-3-5-sonnet-20240620": ChatAnthropic,
+}
+
+
+def load_llm(model: str, api_key: str, temperature: float = 1.0, max_length: int = 2048):
+    model_name = llm_model_translation.get(model)
+    llm_class = llm_classes.get(model_name)
+    if not llm_class:
+        raise ValueError(f"Model {model} not supported.")
+    try:
+        llm = llm_class(model_name=model_name, temperature=temperature, max_tokens=max_length)
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        llm = None
+    return llm
+
 
-def create_db_with_langchain(path):
+def create_db_with_langchain(path: list[str], url_content: dict):
     all_docs = []
-    for file in path:
-        loader = PyMuPDFLoader(file)
-        data = loader.load()
-        # split it into chunks
-        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
-        docs = text_splitter.split_documents(data)
-        all_docs.extend(docs)
-
-    # create the open-source embedding function
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
     embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
-
-    # load it into Chroma
+    if path:
+        for file in path:
+            loader = PyMuPDFLoader(file)
+            data = loader.load()
+            # split it into chunks
+            docs = text_splitter.split_documents(data)
+            all_docs.extend(docs)
+
+    if url_content:
+        for url, content in url_content.items():
+            doc = Document(page_content=content, metadata={"source": url})
+            # split it into chunks
+            docs = text_splitter.split_documents([doc])
+            all_docs.extend(docs)
+
+    assert len(all_docs) > 0, "No PDFs or scrapped data provided"
     db = Chroma.from_documents(all_docs, embedding_function)
     return db
 
 
-def generate_rag(text, model, path):
-    print(f"Generating text using RAG for {model}...")
-    llm = rag_llms[model]
-    db = create_db_with_langchain(path)
+def generate_rag(
+    prompt: str,
+    topic: str,
+    model: str,
+    url_content: dict,
+    path: list[str],
+    temperature: float = 1.0,
+    max_length: int = 2048,
+    api_key: str = "",
+    sys_message="",
+):
+    llm = load_llm(model, api_key, temperature, max_length)
+    if llm is None:
+        print("Failed to load LLM. Aborting operation.")
+        return None
+    db = create_db_with_langchain(path, url_content)
     retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 4, "fetch_k": 20})
-    prompt = hub.pull("rlm/rag-prompt")
+    rag_prompt = hub.pull("rlm/rag-prompt")
 
     def format_docs(docs):
         return "\n\n".join(doc.page_content for doc in docs)
 
-    rag_chain = {"context": retriever | format_docs, "question": RunnablePassthrough()} | prompt | llm
-    return rag_chain.invoke(text).content
-
-
-def generate_groq(text, model):
-    completion = groq_client.chat.completions.create(
-        model=model,
-        messages=[
-            {"role": "user", "content": text},
-            {
-                "role": "assistant",
-                "content": "Please follow the instruction and write about the given topic in approximately the given number of words",
-            },
-        ],
-        temperature=temperature,
-        max_tokens=max_tokens,
-        stream=True,
-        stop=None,
-    )
-    response = ""
-    for i, chunk in enumerate(completion):
-        if i != 0:
-            response += chunk.choices[0].delta.content or ""
-    return response
-
-
-def generate_openai(text, model, openai_client):
-    message = [{"role": "user", "content": text}]
-    response = openai_client.chat.completions.create(
-        model=model,
-        messages=message,
-        temperature=temperature,
-        max_tokens=max_tokens,
-    )
-    return response.choices[0].message.content
-
-
-def generate_gemini(text, model, gemini_client):
-    safety_settings = {
-        generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-    }
-    generation_config = {
-        "max_output_tokens": max_tokens,
-        "temperature": temperature,
-    }
-    response = gemini_client.generate_content(
-        [text],
-        generation_config=generation_config,
-        safety_settings=safety_settings,
-        stream=False,
-    )
-    return response.text
-
-
-def generate_claude(text, model, claude_client):
-    response = claude_client.messages.create(
-        model=model,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        system="You are helpful assistant.",
-        messages=[{"role": "user", "content": [{"type": "text", "text": text}]}],
+    docs = retriever.get_relevant_documents(topic)
+    formatted_docs = format_docs(docs)
+    rag_chain = (
+        {"context": lambda _: formatted_docs, "question": RunnablePassthrough()} | rag_prompt | llm | StrOutputParser()
     )
-    return response.content[0].text.strip()
-
-
-def generate(text, model, path, api=None):
-    if path:
-        result = generate_rag(text, model, path)
-        return result
+    return rag_chain.invoke(prompt)
+
+
+def generate_base(
+    prompt: str, topic: str, model: str, temperature: float, max_length: int, api_key: str, sys_message=""
+):
+    llm = load_llm(model, api_key, temperature, max_length)
+    if llm is None:
+        print("Failed to load LLM. Aborting operation.")
+        return None
+    try:
+        output = llm.invoke(prompt).content
+        return output
+    except Exception as e:
+        print(f"An error occurred while running the model: {e}")
+        return None
+
+
+def generate(
+    prompt: str,
+    topic: str,
+    model: str,
+    url_content: dict,
+    path: list[str],
+    temperature: float = 1.0,
+    max_length: int = 2048,
+    api_key: str = "",
+    sys_message="",
+):
+    if path or url_content:
+        return generate_rag(prompt, topic, model, url_content, path, temperature, max_length, api_key, sys_message)
     else:
-        print(f"Generating text for {model}...")
-        if model == "LLaMA 3":
-            return generate_groq(text, "llama3-70b-8192")
-        elif model == "OpenAI GPT 4o Mini":
-            return generate_openai(text, "gpt-4o-mini", openai_client)
-        elif model == "OpenAI GPT 4o":
-            return generate_openai(text, "gpt-4o", openai_client)
-        elif model == "OpenAI GPT 4":
-            return generate_openai(text, "gpt-4-turbo", openai_client)
-        elif model == "Gemini 1.5 Pro":
-            return generate_gemini(text, "", gemini_client)
-        elif model == "Claude Sonnet 3.5":
-            return generate_claude(text, "claude-3-5-sonnet-20240620", claude_client)
+        return generate_base(prompt, topic, model, temperature, max_length, api_key, sys_message)

app.py

@@ -16,7 +16,7 @@ from transformers import GPT2LMHeadModel, GPT2TokenizerFast
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
 from utils import remove_special_characters
-from plagiarism import google_search, months, domain_list, build_date
+from google_search import google_search, months, domain_list, build_date
 from humanize import paraphrase_text, device
 from ai_generate import generate
 
@@ -261,6 +261,8 @@ def generate_prompt(settings: Dict[str, str]) -> str:
     prompt = f"""
     I am a {settings['role']}
     Write a {settings['article_length']} words (around) {settings['format']} on {settings['topic']}.
+    Context:
+    - {settings['context']}
     
     Style and Tone:
     - Writing style: {settings['writing_style']}
@@ -282,8 +284,6 @@ def generate_prompt(settings: Dict[str, str]) -> str:
     - Add a "References" section in the format "References:" on a new line at the end with at least 3 credible detailed sources, formatted as [1], [2], etc. with each source on their own line
     - Do not repeat sources
     - Do not make any headline, title bold.
-
-    {settings['sources']}
     
     Ensure proper paragraph breaks for better readability.
     Avoid any references to artificial intelligence, language models, or the fact that this is generated by an AI, and do not mention something like here is the article etc.
@@ -296,15 +296,16 @@ def regenerate_prompt(settings: Dict[str, str]) -> str:
     I am a {settings['role']}
     "{settings['generated_article']}"
     Edit the given text based on user comments.
+    User Comments:
+    - {settings['user_comments']}
     
-    Comments:
+    Requirements:
     - Don't start with "Here is a...", start with the requested text directly
-    - {settings['user_comments']}
     - The original content should not be changed. Make minor modifications based on user comments above.
     - Keep the references the same as the given text in the same format.
     - Do not make any headline, title bold.
-
-    {settings['sources']}
+    Context:
+    - {settings['context']}
     
     Ensure proper paragraph breaks for better readability.
     Avoid any references to artificial intelligence, language models, or the fact that this is generated by an AI, and do not mention something like here is the article etc.
@@ -315,6 +316,7 @@ def regenerate_prompt(settings: Dict[str, str]) -> str:
 def generate_article(
     input_role: str,
     topic: str,
+    context: str,
     keywords: str,
     article_length: str,
     format: str,
@@ -328,14 +330,16 @@ def generate_article(
     conclusion_type: str,
     ai_model: str,
     content_string: str,
+    url_content: str = None,
     # api_key: str = None,
-    pdf_file_input=None,
+    pdf_file_input: list[str] = None,
     generated_article: str = None,
     user_comments: str = None,
 ) -> str:
     settings = {
         "role": input_role,
         "topic": topic,
+        "context": context,
         "keywords": [k.strip() for k in keywords.split(",")],
         "article_length": article_length,
         "format": format,
@@ -360,7 +364,9 @@ def generate_article(
     print("Generated Prompt...\n", prompt)
     article = generate(
         prompt,
+        topic,
         ai_model,
+        url_content,
         pdf_file_input,  # api_key
     )
 
@@ -399,6 +405,7 @@ def update_visibility_api(model: str):
 def generate_and_format(
     input_role,
     topic,
+    context,
     keywords,
     article_length,
     format,
@@ -410,7 +417,7 @@ def generate_and_format(
     references,
     num_examples,
     conclusion_type,
-    ai_model,
+    # ai_model,
     # api_key,
     google_search_check,
     year_from,
@@ -428,6 +435,7 @@ def generate_and_format(
 ):
     content_string = ""
     url_content = None
+    ai_model = "Claude Sonnet 3.5"
     if google_search_check:
         date_from = build_date(year_from, month_from, day_from)
         date_to = build_date(year_to, month_to, day_to)
@@ -450,6 +458,7 @@ def generate_and_format(
     article = generate_article(
         input_role,
         topic,
+        context,
         keywords,
         article_length,
         format,
@@ -463,6 +472,7 @@ def generate_and_format(
         conclusion_type,
         ai_model,
         content_string,
+        url_content,
         # api_key,
         pdf_file_input,
         generated_article,
@@ -500,6 +510,11 @@ def create_interface():
                         placeholder="Enter the main topic of your article",
                         elem_classes="input-highlight-pink",
                     )
+                    input_context = gr.Textbox(
+                        label="Context",
+                        placeholder="Provide some context for your topic",
+                        elem_classes="input-highlight-pink",
+                    )
                     input_keywords = gr.Textbox(
                         label="Keywords",
                         placeholder="Enter comma-separated keywords",
@@ -667,23 +682,24 @@ def create_interface():
                     gr.Markdown("# Add Optional PDF Files with Information", elem_classes="text-center text-3xl mb-6")
                     pdf_file_input = gr.File(label="Upload PDF(s)", file_count="multiple", file_types=[".pdf"])
 
-                with gr.Group():
-                    gr.Markdown("## AI Model Configuration", elem_classes="text-xl mb-4")
-                    ai_generator = gr.Dropdown(
-                        choices=[
-                            "OpenAI GPT 4",
-                            "OpenAI GPT 4o",
-                            "OpenAI GPT 4o Mini",
-                            "Claude Sonnet 3.5",
-                            "Gemini 1.5 Pro",
-                            "LLaMA 3",
-                        ],
-                        value="OpenAI GPT 4o Mini",
-                        label="AI Model",
-                        elem_classes="input-highlight-pink",
-                    )
-                    # input_api = gr.Textbox(label="API Key", visible=False)
-                    # ai_generator.change(update_visibility_api, ai_generator, input_api)
+                # HIDE AI MODEL SELECTION
+                # with gr.Group():
+                #     gr.Markdown("## AI Model Configuration", elem_classes="text-xl mb-4")
+                #     ai_generator = gr.Dropdown(
+                #         choices=[
+                #             "OpenAI GPT 4",
+                #             "OpenAI GPT 4o",
+                #             "OpenAI GPT 4o Mini",
+                #             "Claude Sonnet 3.5",
+                #             "Gemini 1.5 Pro",
+                #             "LLaMA 3",
+                #         ],
+                #         value="OpenAI GPT 4o Mini",
+                #         label="AI Model",
+                #         elem_classes="input-highlight-pink",
+                #     )
+                # input_api = gr.Textbox(label="API Key", visible=False)
+                # ai_generator.change(update_visibility_api, ai_generator, input_api)
 
                 generate_btn = gr.Button("Generate Article", variant="primary")
 
@@ -762,6 +778,7 @@ def create_interface():
             inputs=[
                 input_role,
                 input_topic,
+                input_context,
                 input_keywords,
                 input_length,
                 input_format,
@@ -773,7 +790,7 @@ def create_interface():
                 input_references,
                 input_num_examples,
                 input_conclusion,
-                ai_generator,
+                # ai_generator,
                 # input_api,
                 google_search_check,
                 year_from,
@@ -795,6 +812,7 @@ def create_interface():
             inputs=[
                 input_role,
                 input_topic,
+                input_context,
                 input_keywords,
                 input_length,
                 input_format,
@@ -806,7 +824,7 @@ def create_interface():
                 input_references,
                 input_num_examples,
                 input_conclusion,
-                ai_generator,
+                # ai_generator,
                 # input_api,
                 google_search_check,
                 year_from,